"""실험: Set Design 프로토타입

1단계: LLM이 location 설명 + 해당 샷 목록을 보고 세트 구도 N개 제안
2단계: 세트 구도별 배경 이미지 생성 (인물 점선 실루엣 포함)
3단계: LLM이 전체 T2I 프롬프트 + 세트 구도 설명을 보고 샷별 배정

Usage:
    cd backend && .venv/bin/python -m scripts.experiment_set_design
"""

import json
import pathlib
import sys
import time

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))
import os; os.environ.setdefault("RUNNING_SCRIPT", "1")

from app.core.config import settings
from app.modules.llm.gemini_image_client import GeminiImageClient
from app.modules.llm.llm_client import call_text

# ── 프로젝트 데이터 경로 ──
BASE = pathlib.Path(
    "/Users/manta/Documents/Projects/TheRoad-I1/projects/"
    "b789d6ce-f474-4f49-9388-b03c9d95020e/checkpoints/episodes/"
    "0a4d9c09-099f-4eee-b8e2-ca0170a9431d"
)
OUT_DIR = pathlib.Path(__file__).resolve().parent.parent.parent / "frontend" / "public" / "experiment" / "set_chain"


# ──────────────────────────────────────────────
# 데이터 로드
# ──────────────────────────────────────────────

def load_location_info(short_id: str) -> dict:
    """entity_t2i manifest에서 location 정보 로드."""
    with open(BASE / "entity_t2i" / "manifest.json") as f:
        data = json.load(f)
    for name, info in data.get("data", data).get("completed", {}).items():
        if info.get("short_id") == short_id:
            return info
    raise ValueError(f"Location {short_id} not found")


def load_shots_for_location(short_id: str) -> list[dict]:
    """scene_detail에서 해당 location을 사용하는 샷 목록 로드."""
    with open(BASE / "scene_detail" / "manifest.json") as f:
        data = json.load(f)

    shots = []
    for s in data["data"]["scenes"]:
        ve = s.get("visible_entities", [])
        if short_id not in ve:
            continue
        chars = [e for e in ve if e.startswith("C")]
        t2i_list = s.get("t2i_variations", [])
        t2i_prompt = t2i_list[0].get("t2i_prompt", "") if t2i_list else ""
        shots.append({
            "shot_label": f"S{s['scene_index']:02d}-Shot{s.get('_shot_index', 1)}",
            "scene_index": s["scene_index"],
            "shot_index": s.get("_shot_index", 1),
            "beat_title": s.get("beat_title", ""),
            "representative_moment": s.get("representative_moment", ""),
            "characters": chars,
            "character_count": len(chars),
            "t2i_prompt": t2i_prompt,
        })
    return shots


def load_character_names() -> dict:
    """entity_merge에서 캐릭터 이름 로드."""
    with open(BASE / "entity_merge" / "manifest.json") as f:
        data = json.load(f)
    md = data.get("data", data)
    names = {}
    for c in md.get("characters", []):
        sid = c.get("short_id", "")
        if sid:
            names[sid] = c.get("name", sid)
    return names


# ──────────────────────────────────────────────
# 1단계: LLM이 세트 구도 제안
# ──────────────────────────────────────────────

STEP1_SYSTEM = """You are a professional film production designer and cinematographer.
Your job is to design camera-ready SET COMPOSITIONS for a given location
that will cover all the shots planned for that location.

Each set composition is a specific camera angle/position/framing of the same physical space.
Think of it as: "if we built this set on a soundstage, what are the different angles
we'd pre-plan to shoot from?"

Rules:
- Each composition must be visually distinct (different angle, height, depth, framing)
- Consider the dramatic needs of the shots (intimate vs wide, tense vs calm)
- The number of compositions should be efficient: enough to cover all shot needs,
  but not wasteful. Minimum 2, but scale with shot count and variety.
- Output valid JSON only."""

STEP1_USER_TEMPLATE = """## Location: {loc_name} ({loc_id})
{loc_description}

Visual traits: {visual_traits}

## Shots using this location ({shot_count} shots):
{shots_text}

## Task
Design set compositions (camera angles/positions/framings) for this location.
Consider the shots above — their action, mood, character count, and dramatic needs.
Each composition should serve multiple shots where possible.

Respond in JSON:
```json
{{
  "compositions": [
    {{
      "id": "SET_A",
      "name_ko": "한국어 구도 이름",
      "name_en": "English composition name",
      "description": "Detailed description of camera position, angle, height, what's visible, lighting direction, mood",
      "suited_for": "What kind of shots this is good for (mood, action type, character count)",
      "t2i_prompt": "Full T2I prompt for generating this background. Photorealistic cinematic. Include the location details. NO people — instead include exactly 3 person-shaped DOTTED LINE silhouettes (dashed outlines only, no fill) at different heights and depths as actor placement guides."
    }}
  ],
  "reasoning": "Brief explanation of why you chose these compositions and how they cover the shots"
}}
```"""


def step1_propose_compositions(location: dict, shots: list[dict], char_names: dict) -> dict:
    """LLM에게 세트 구도 제안 요청."""
    shots_text = ""
    for s in shots:
        char_str = ", ".join(f"{c}({char_names.get(c, '?')})" for c in s["characters"])
        shots_text += (
            f"- **{s['shot_label']}** | {s['character_count']}명 [{char_str}]\n"
            f"  Beat: {s['beat_title']}\n"
            f"  Moment: {s['representative_moment'][:150]}\n\n"
        )

    user_prompt = STEP1_USER_TEMPLATE.format(
        loc_name=location["name"],
        loc_id=location["short_id"],
        loc_description=location["description"],
        visual_traits=", ".join(location.get("visual_traits", [])),
        shot_count=len(shots),
        shots_text=shots_text,
    )

    print("[Step 1] Asking LLM to propose set compositions...")
    result = call_text(
        system_prompt=STEP1_SYSTEM,
        user_prompt=user_prompt,
        step="set_design_propose",
    )

    # Parse JSON from response
    text = result if isinstance(result, str) else result.get("text", "")
    # Strip markdown code fences if present
    if "```json" in text:
        text = text.split("```json")[1].split("```")[0]
    elif "```" in text:
        text = text.split("```")[1].split("```")[0]

    return json.loads(text)


# ──────────────────────────────────────────────
# 2단계: 세트 구도별 배경 이미지 생성
# ──────────────────────────────────────────────

def step2_generate_set_images(compositions: list[dict]) -> dict[str, pathlib.Path]:
    """체이닝 방식으로 세트 이미지 생성.

    A: 참조 0개 (기준)
    B: A 참조
    C: A+B 참조
    D: A+B+C 참조
    """
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    client = GeminiImageClient(model=settings.gemini_image_model)
    client.set_context(step="set_design_generate", operation_type="experiment")

    results = {}
    prev_images: list[bytes] = []  # 이전 생성 이미지 누적

    for i, comp in enumerate(compositions):
        comp_id = comp["id"]
        prompt = comp["t2i_prompt"]

        # 체이닝: 이전 이미지들을 참조로 추가
        ref_label = f"(refs: {len(prev_images)})" if prev_images else "(no refs)"
        chain_prompt = prompt
        if prev_images:
            chain_prompt = (
                f"{prompt}\n\n"
                f"IMPORTANT: This is the SAME physical room/location as the reference image(s) below. "
                f"Maintain the same wall colors, flooring, furniture style, window shape, and overall condition. "
                f"Only the camera angle and framing should differ."
            )

        out_path = OUT_DIR / f"{comp_id}.png"
        print(f"[Step 2] Generating {comp_id} ({comp['name_ko']}) {ref_label}... ", end="", flush=True)
        t0 = time.time()
        try:
            labeled_refs = None
            if prev_images:
                labeled_refs = [
                    (f"Same room, different angle (ref {j+1}):", img)
                    for j, img in enumerate(prev_images)
                ]

            img_bytes, resp_ms = client.generate_image(
                prompt=chain_prompt,
                aspect_ratio="16:9",
                labeled_references=labeled_refs,
            )
            out_path.write_bytes(img_bytes)
            elapsed = time.time() - t0
            print(f"OK  {len(img_bytes)//1024}KB  {elapsed:.1f}s")
            results[comp_id] = out_path
            prev_images.append(img_bytes)
        except Exception as e:
            print(f"FAIL: {e}")

    return results


# ──────────────────────────────────────────────
# 3단계: LLM이 전체 T2I 보고 샷별 세트 배정
# ──────────────────────────────────────────────

STEP3_SYSTEM = """You are a film director assigning pre-designed set compositions to individual shots.
You must consider:
1. How well the composition's angle/mood matches the shot's dramatic needs
2. Visual VARIETY — avoid assigning the same composition to consecutive shots in the same scene
3. Character blocking — does the composition's space work for the number of characters?
4. Story flow — camera angle changes between shots should feel intentional

Output valid JSON only."""

STEP3_USER_TEMPLATE = """## Available Set Compositions for "{loc_name}":
{compositions_text}

## Shots to assign (each needs exactly ONE composition):
{shots_text}

## Task
Assign each shot to the best-fitting composition.
Ensure visual variety — consecutive shots in the same scene should use different compositions when possible.

Respond in JSON:
```json
{{
  "assignments": [
    {{
      "shot_label": "S05-Shot1",
      "composition_id": "SET_A",
      "reason": "Brief reason for this choice"
    }}
  ]
}}
```"""


def step3_assign_compositions(
    location: dict, compositions: list[dict], shots: list[dict], char_names: dict
) -> dict:
    """LLM이 전체 T2I 프롬프트 기반으로 샷별 세트 배정."""
    compositions_text = ""
    for comp in compositions:
        compositions_text += (
            f"### {comp['id']}: {comp['name_ko']} ({comp['name_en']})\n"
            f"  {comp['description']}\n"
            f"  Suited for: {comp['suited_for']}\n\n"
        )

    shots_text = ""
    for s in shots:
        char_str = ", ".join(f"{c}({char_names.get(c, '?')})" for c in s["characters"])
        shots_text += (
            f"- **{s['shot_label']}** | {s['character_count']}명 [{char_str}]\n"
            f"  Beat: {s['beat_title']}\n"
            f"  Moment: {s['representative_moment'][:200]}\n"
            f"  T2I: {s['t2i_prompt'][:200]}\n\n"
        )

    user_prompt = STEP3_USER_TEMPLATE.format(
        loc_name=location["name"],
        compositions_text=compositions_text,
        shots_text=shots_text,
    )

    print("[Step 3] Asking LLM to assign compositions to shots...")
    result = call_text(
        system_prompt=STEP3_SYSTEM,
        user_prompt=user_prompt,
        step="set_design_assign",
    )

    text = result if isinstance(result, str) else result.get("text", "")
    if "```json" in text:
        text = text.split("```json")[1].split("```")[0]
    elif "```" in text:
        text = text.split("```")[1].split("```")[0]

    return json.loads(text)


# ──────────────────────────────────────────────
# Main
# ──────────────────────────────────────────────

def main():
    TARGET_LOC = "L05"

    print(f"=== Set Design Prototype — {TARGET_LOC} ===\n")

    # Load data
    location = load_location_info(TARGET_LOC)
    shots = load_shots_for_location(TARGET_LOC)
    char_names = load_character_names()

    print(f"Location: {location['name']} ({TARGET_LOC})")
    print(f"Shots: {len(shots)}개")
    print(f"Characters: {', '.join(f'{k}={v}' for k, v in char_names.items() if k.startswith('C'))}")
    print()

    # Step 1: LLM proposes set compositions
    step1_result = step1_propose_compositions(location, shots, char_names)
    compositions = step1_result["compositions"]
    print(f"\n[Step 1 결과] {len(compositions)}개 세트 구도 제안됨:")
    for comp in compositions:
        print(f"  {comp['id']}: {comp['name_ko']} — {comp['suited_for'][:80]}")
    print(f"  Reasoning: {step1_result.get('reasoning', '')[:200]}")
    print()

    # Save step1 result
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    with open(OUT_DIR / "step1_compositions.json", "w") as f:
        json.dump(step1_result, f, ensure_ascii=False, indent=2)

    # Step 2: Generate set images
    print()
    image_paths = step2_generate_set_images(compositions)
    print(f"\n[Step 2 결과] {len(image_paths)}개 이미지 생성됨")
    print()

    # Step 3: Assign compositions to shots
    step3_result = step3_assign_compositions(location, compositions, shots, char_names)
    assignments = step3_result["assignments"]
    print(f"\n[Step 3 결과] 샷별 세트 배정:")
    for a in assignments:
        comp_id = a["composition_id"]
        comp_name = next((c["name_ko"] for c in compositions if c["id"] == comp_id), "?")
        print(f"  {a['shot_label']} → {comp_id} ({comp_name})")
        print(f"    이유: {a['reason'][:100]}")
    print()

    # Save step3 result
    with open(OUT_DIR / "step3_assignments.json", "w") as f:
        json.dump(step3_result, f, ensure_ascii=False, indent=2)

    # Summary
    print("=== Summary ===")
    from collections import Counter
    usage = Counter(a["composition_id"] for a in assignments)
    for comp in compositions:
        cid = comp["id"]
        count = usage.get(cid, 0)
        img = f"✓ {image_paths[cid].name}" if cid in image_paths else "✗ no image"
        print(f"  {cid} ({comp['name_ko']}): {count}개 샷 배정 [{img}]")

    print(f"\nFiles saved to: {OUT_DIR}")
    print("Images:")
    for cid, path in image_paths.items():
        print(f"  http://192.168.231.91:3002/experiment/set_design/{path.name}")


if __name__ == "__main__":
    main()