"""실험 v3: 2-pass 세트 디자인 (parent→child→shot)

Phase 1 (LLM 분석): gemini-pro가 전체 T2I를 보고
  - 기본 세트 구도 제안 (parent)
  - 상태별 변형 세트 제안 (child) + 배경 프롬프트
  - 샷별 T2I를 인물/액션만 남기도록 수정 제안

Phase 2 (이미지 생성):
  - parent 세트 생성 (체이닝)
  - child 세트 생성 (parent 참조)
  - 샷 이미지 생성 (child + 인물 refs)

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

import json
import pathlib
import re
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.core.database import SessionLocal
from app.models.project import ImageAsset, EntityCanon as Entity
from app.modules.llm.gemini_image_client import GeminiImageClient
from app.modules.llm.llm_client import call_text

PROJECT_ID = "b789d6ce-f474-4f49-9388-b03c9d95020e"
EPISODE_ID = "0a4d9c09-099f-4eee-b8e2-ca0170a9431d"
BASE = pathlib.Path(
    f"/Users/manta/Documents/Projects/TheRoad-I1/projects/{PROJECT_ID}"
    f"/checkpoints/episodes/{EPISODE_ID}"
)
EXP = pathlib.Path(__file__).resolve().parent.parent.parent / "frontend" / "public" / "experiment"
OUT_DIR = EXP / "set_v3"


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

def load_l05_data():
    with open(BASE / "scene_detail" / "manifest.json") as f:
        data = json.load(f)

    with open(BASE / "entity_merge" / "manifest.json") as f:
        merge = json.load(f).get("data", {})

    with open(BASE / "entity_t2i" / "manifest.json") as f:
        t2i_ent = json.load(f).get("data", {}).get("completed", {})

    # Names
    names = {}
    for cat in ["characters", "locations", "props"]:
        for e in merge.get(cat, []):
            sid = e.get("short_id", "")
            if sid:
                names[sid] = e.get("name", sid)
    for n, info in t2i_ent.items():
        sid = info.get("short_id", "")
        if sid and sid.startswith("O"):
            names[sid] = info.get("name", n)

    # Location info
    loc_info = None
    for n, info in t2i_ent.items():
        if info.get("short_id") == "L05":
            loc_info = info
            break

    # Shots
    shots = []
    for s in data["data"]["scenes"]:
        ve = s.get("visible_entities", [])
        if "L05" not in ve:
            continue
        si = s["scene_index"]
        shi = s.get("_shot_index", 1)
        t2i_list = s.get("t2i_variations", [])
        t2i_prompt = t2i_list[0].get("t2i_prompt", "") if t2i_list else ""

        combos = list(dict.fromkeys(re.findall(r'C\d{2,3}O\d{2,3}', t2i_prompt)))

        shots.append({
            "label": f"S{si:02d}-Shot{shi}",
            "scene_index": si,
            "shot_index": shi,
            "beat_title": s.get("beat_title", ""),
            "representative_moment": s.get("representative_moment", ""),
            "visible_entities": ve,
            "character_combos": combos,
            "t2i_prompt": t2i_prompt,
        })

    # Entity text map
    etm = {}
    for n, info in t2i_ent.items():
        sid = info.get("short_id", "")
        desc = info.get("short_description", info.get("description", ""))
        if sid and desc:
            etm[sid] = desc

    return loc_info, shots, names, etm


# ──────────────────────────────────────────────
# Phase 1: LLM 분석 (gemini-pro)
# ──────────────────────────────────────────────

PHASE1_SYSTEM = """You are a professional film production designer analyzing shots for set design.
All shots take place in the SAME physical location. Your job:

1. Identify the BASE SET COMPOSITIONS (camera angles) needed.
   These are the "parent" sets — the room in its default state, empty, no people.

2. Identify STATE VARIANTS for each composition.
   States = story-driven changes to the room (blood, mess, unnaturally clean, etc.)
   Each state variant is a "child" of a parent composition.
   If a shot uses the room in its default state, no child is needed (use parent directly).

3. For each shot, SPLIT the T2I prompt into:
   - background_prompt: ONLY room state changes (blood, mess, etc.) — used to generate the child set
   - character_prompt: ONLY character actions, poses, expressions, props they HOLD — used for final image generation
   - The character_prompt must NOT contain any room/background descriptions since the background reference image already provides that.

Rules:
- Parent compositions: empty room, no people, no state changes. Different camera angles only.
- Child compositions: parent + state changes only. Still no people.
- character_prompt: Remove ALL background/room descriptions. Keep only: character appearance, action, pose, expression, held props, camera framing, lighting mood.
- character_prompt: NEVER include skin tone / face color modifiers (pale, drained, flushed, ashen, gray, white face, etc.) — the image model renders these literally and destroys the face. Express fear/shock through body language and expression words only (frozen, trembling, wide eyes, clenched jaw, etc.)
- If a shot doesn't need state changes (room is normal), set child_id = null (use parent directly).
- Output valid JSON only.
- Respond in Korean for names, English for prompts."""

PHASE1_USER = """## Location: {loc_name} (L05)
{loc_description}
Visual traits: {visual_traits}

## All shots and their FULL T2I prompts:
{shots_text}

## Output JSON format:
```json
{{
  "parent_sets": [
    {{
      "id": "SET_A",
      "name": "구도 이름",
      "description": "Camera angle/position description",
      "t2i_prompt": "Photorealistic cinematic shot of... (EMPTY room, NO people, NO state changes, default condition)"
    }}
  ],
  "child_sets": [
    {{
      "id": "SET_A_bloody",
      "parent_id": "SET_A",
      "state_name": "혈흔 상태",
      "state_description": "What changed from parent",
      "t2i_prompt": "Same room as reference image, but with: blood pooled on floor, bloody footprints... (NO people)"
    }}
  ],
  "shot_assignments": [
    {{
      "shot_label": "S12-Shot1",
      "parent_id": "SET_B",
      "child_id": "SET_B_bloody",
      "character_prompt": "Cleaned T2I — character actions/poses ONLY, no background descriptions"
    }}
  ]
}}
```"""


def phase1_analyze(loc_info, shots, names):
    shots_text = ""
    for s in shots:
        chars = ", ".join(f"{c}({names.get(c[:3], '?')}+{names.get(c[3:], '?')})" for c in s["character_combos"])
        shots_text += (
            f"### {s['label']} | Characters: {chars}\n"
            f"Beat: {s['beat_title']}\n"
            f"Moment: {s['representative_moment'][:200]}\n"
            f"T2I:\n```\n{s['t2i_prompt']}\n```\n\n"
        )

    user_prompt = PHASE1_USER.format(
        loc_name=loc_info["name"],
        loc_description=loc_info["description"],
        visual_traits=", ".join(loc_info.get("visual_traits", [])),
        shots_text=shots_text,
    )

    print("[Phase 1] Gemini Pro 분석 중...", flush=True)
    result = call_text(
        step="scene_director",  # gemini-pro mapped step
        system_prompt=PHASE1_SYSTEM,
        user_prompt=user_prompt,
        temperature=0.1,
    )

    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)


# ──────────────────────────────────────────────
# Phase 2: 이미지 생성
# ──────────────────────────────────────────────

def phase2_generate_parents(parent_sets):
    """Parent 세트 생성 (체이닝)."""
    parent_dir = OUT_DIR / "parents"
    parent_dir.mkdir(parents=True, exist_ok=True)
    client = GeminiImageClient(model=settings.gemini_image_model)
    client.set_context(step="set_v3_parent", operation_type="experiment")

    results = {}
    prev_images = []

    for p in parent_sets:
        pid = p["id"]
        prompt = p["t2i_prompt"]
        if prev_images:
            prompt += (
                "\n\nIMPORTANT: This is the SAME physical room as the reference image(s). "
                "Maintain identical wall colors, flooring, furniture style, window shape. "
                "Only the camera angle differs."
            )

        ref_label = f"(refs: {len(prev_images)})" if prev_images else "(no refs)"
        print(f"  [parent] {pid} ({p['name']}) {ref_label}... ", end="", flush=True)

        labeled_refs = [
            (f"Same room ref {j+1}:", img) for j, img in enumerate(prev_images)
        ] if prev_images else None

        t0 = time.time()
        try:
            img_bytes, _ = client.generate_image(prompt=prompt, aspect_ratio="16:9", labeled_references=labeled_refs)
            out = parent_dir / f"{pid}.png"
            out.write_bytes(img_bytes)
            print(f"OK {len(img_bytes)//1024}KB {time.time()-t0:.1f}s")
            results[pid] = {"path": out, "bytes": img_bytes}
            prev_images.append(img_bytes)
        except Exception as e:
            print(f"FAIL: {e}")

    return results


def phase2_generate_children(child_sets, parent_results):
    """Child 세트 생성 (parent 참조)."""
    child_dir = OUT_DIR / "children"
    child_dir.mkdir(parents=True, exist_ok=True)
    client = GeminiImageClient(model=settings.gemini_image_model)
    client.set_context(step="set_v3_child", operation_type="experiment")

    results = {}
    for c in child_sets:
        cid = c["id"]
        parent_id = c["parent_id"]
        parent = parent_results.get(parent_id)
        if not parent:
            print(f"  [child] {cid}: SKIP (no parent {parent_id})")
            continue

        prompt = c["t2i_prompt"]
        prompt += (
            "\n\nCRITICAL: Use the reference image as the base room. "
            "Keep the exact same walls, floor, furniture layout, window. "
            "Only apply the described state changes. NO people."
        )

        labeled_refs = [("Base room — maintain this exact room structure:", parent["bytes"])]

        print(f"  [child] {cid} ({c['state_name']}) ← {parent_id}... ", end="", flush=True)
        t0 = time.time()
        try:
            img_bytes, _ = client.generate_image(prompt=prompt, aspect_ratio="16:9", labeled_references=labeled_refs)
            out = child_dir / f"{cid}.png"
            out.write_bytes(img_bytes)
            print(f"OK {len(img_bytes)//1024}KB {time.time()-t0:.1f}s")
            results[cid] = {"path": out, "bytes": img_bytes}
        except Exception as e:
            print(f"FAIL: {e}")

    return results


def get_character_refs(db, shot, t2i_text):
    """DB에서 인물/소품 참조 이미지 로드 (composite 우선)."""
    ve = shot["visible_entities"]
    labeled_refs = []

    for sid in ve:
        if sid.startswith("L"):
            continue
        entity = db.query(Entity).filter(Entity.project_id == PROJECT_ID, Entity.short_id == sid).first()
        if not entity:
            continue
        eid = entity.id
        etype = entity.entity_type

        if etype == "character":
            # character_prompt에 C## 언급 없으면 스킵
            if not re.search(rf'(?<![CO\d]){sid}', t2i_text):
                continue
            m = re.search(rf'{sid}(O\d{{2,3}})', t2i_text)
            if m:
                outfit_sid = m.group(1)
                outlook = db.query(Entity).filter(Entity.project_id == PROJECT_ID, Entity.short_id == outfit_sid).first()
                if outlook:
                    comp = db.query(ImageAsset).filter(
                        ImageAsset.project_id == PROJECT_ID,
                        ImageAsset.asset_type == "reference",
                        ImageAsset.prompt_used.like(f"%composite:{eid}:{outlook.id}%"),
                    ).order_by(ImageAsset.is_primary.desc(), ImageAsset.created_at.desc()).first()
                    if comp and pathlib.Path(comp.file_path).exists():
                        labeled_refs.append((f"{sid}{outfit_sid} — character wearing outfit", pathlib.Path(comp.file_path).read_bytes()))
                        continue

            char_ref = db.query(ImageAsset).filter(
                ImageAsset.project_id == PROJECT_ID, ImageAsset.entity_id == eid,
                ImageAsset.asset_type == "reference", ImageAsset.is_primary == 1,
            ).first()
            if char_ref and pathlib.Path(char_ref.file_path).exists():
                labeled_refs.append((f"{sid} — character face identity", pathlib.Path(char_ref.file_path).read_bytes()))

        elif etype == "prop":
            # prop은 t2i_text에 실제 언급된 경우만 참조 추가
            if not re.search(rf'(?<![A-Z]){sid}(?!\d)', t2i_text):
                continue
            prop_ref = db.query(ImageAsset).filter(
                ImageAsset.project_id == PROJECT_ID, ImageAsset.entity_id == eid,
                ImageAsset.asset_type == "reference", ImageAsset.is_primary == 1,
            ).first()
            if prop_ref and pathlib.Path(prop_ref.file_path).exists():
                labeled_refs.append((f"{sid} — prop/object", pathlib.Path(prop_ref.file_path).read_bytes()))

    return labeled_refs


def phase2_generate_shots(assignments, parent_results, child_results, shots_data, etm):
    """최종 샷 이미지 생성."""
    shot_dir = OUT_DIR / "shots"
    shot_dir.mkdir(parents=True, exist_ok=True)
    client = GeminiImageClient(model=settings.gemini_image_model)
    client.set_context(step="set_v3_shot", operation_type="experiment",
                       project_id=PROJECT_ID, episode_id=EPISODE_ID)

    db = SessionLocal()
    try:
        for a in assignments:
            label = a["shot_label"]
            parent_id = a["parent_id"]
            child_id = a.get("child_id")
            char_prompt = a["character_prompt"]

            # 배경 이미지 선택: child가 있으면 child, 없으면 parent
            if child_id and child_id in child_results:
                bg = child_results[child_id]
                bg_label = child_id
            elif parent_id in parent_results:
                bg = parent_results[parent_id]
                bg_label = parent_id
            else:
                print(f"  [shot] {label}: SKIP (no bg)")
                continue

            # 원본 shot 데이터에서 VE 가져오기
            shot = next((s for s in shots_data if s["label"] == label), None)
            if not shot:
                continue

            # 인물/소품 참조 — character_prompt에 언급된 것만
            # character_prompt 기준으로 C##O##/P## 매칭 (VE 전체가 아닌 실제 사용분만)
            entity_refs = get_character_refs(db, shot, char_prompt)

            # 배경을 첫 번째 참조로
            labeled_refs = [("SET BACKGROUND — use this exact room as the background", bg["bytes"])]
            labeled_refs.extend(entity_refs)

            # 최종 프롬프트 빌드 (character_prompt 사용 — 배경 묘사 제거된 버전)
            from app.services.prompt_service import build_final_scene_prompt as _build_final_scene_prompt
            final_prompt = _build_final_scene_prompt(
                t2i_prompt=char_prompt,
                labeled_refs=labeled_refs,
                style_context="",
                scene_index=shot["scene_index"],
                entity_text_map=etm,
            )

            print(f"  [shot] {label} → {bg_label} | refs: 1+{len(entity_refs)} | ", end="", flush=True)
            t0 = time.time()
            try:
                img_bytes, _ = client.generate_image(
                    prompt=final_prompt, aspect_ratio="16:9", labeled_references=labeled_refs,
                )
                filename = label.replace("-", "_") + ".png"
                (shot_dir / filename).write_bytes(img_bytes)
                print(f"OK {len(img_bytes)//1024}KB {time.time()-t0:.1f}s")
            except Exception as e:
                print(f"FAIL ({time.time()-t0:.1f}s): {e}")
    finally:
        db.close()


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

def main():
    OUT_DIR.mkdir(parents=True, exist_ok=True)

    loc_info, shots, names, etm = load_l05_data()
    print(f"=== Set Design v3 — L05 ({loc_info['name']}) ===")
    print(f"Shots: {len(shots)}")
    print()

    # Phase 1: LLM Analysis
    print("=" * 50)
    print("Phase 1: Gemini Pro 분석")
    print("=" * 50)
    analysis = phase1_analyze(loc_info, shots, names)

    parents = analysis["parent_sets"]
    children = analysis["child_sets"]
    assignments = analysis["shot_assignments"]

    print(f"\n  Parent 세트: {len(parents)}개")
    for p in parents:
        print(f"    {p['id']}: {p['name']}")
    print(f"  Child 세트: {len(children)}개")
    for c in children:
        print(f"    {c['id']} ← {c['parent_id']}: {c['state_name']}")
    print(f"  샷 배정: {len(assignments)}개")
    for a in assignments:
        bg = a.get("child_id") or a["parent_id"]
        print(f"    {a['shot_label']} → {bg}")

    # Save analysis
    with open(OUT_DIR / "phase1_analysis.json", "w") as f:
        json.dump(analysis, f, ensure_ascii=False, indent=2)
    print(f"\n  분석 저장: {OUT_DIR / 'phase1_analysis.json'}")

    # Phase 2: Image Generation
    print()
    print("=" * 50)
    print("Phase 2: 이미지 생성")
    print("=" * 50)

    print("\n[2-1] Parent 세트 생성 (체이닝)")
    parent_results = phase2_generate_parents(parents)

    print("\n[2-2] Child 세트 생성 (parent 참조)")
    child_results = phase2_generate_children(children, parent_results)

    print("\n[2-3] 샷 이미지 생성")
    phase2_generate_shots(assignments, parent_results, child_results, shots, etm)

    # Summary
    print()
    print("=" * 50)
    print("URLs")
    print("=" * 50)
    base_url = "http://192.168.231.91:3002/experiment/set_v3"
    print("\nParent 세트:")
    for p in sorted((OUT_DIR / "parents").glob("*.png")):
        print(f"  {base_url}/parents/{p.name}")
    print("\nChild 세트:")
    for p in sorted((OUT_DIR / "children").glob("*.png")):
        print(f"  {base_url}/children/{p.name}")
    print("\n샷:")
    for p in sorted((OUT_DIR / "shots").glob("*.png")):
        print(f"  {base_url}/shots/{p.name}")


if __name__ == "__main__":
    main()