"""실험 v4: 세트 디자인 최종 프로토타입

변경사항 (v3 대비):
- parent 최대 4개 (같은 공간 최대 2구도)
- 생성 시마다 VLM blueprint 체이닝 검증
- 참조 이미지 최대 2장 (LLM 선택)
- character_prompt에서 피부색 묘사 제거
- prop은 character_prompt에 언급된 것만 참조
- gore 표현 완화 (SAFETY 방지)

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

import json
import pathlib
import re
import sys
import time
import base64

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, call_structured

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_v4"


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

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 = {}
    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)

    loc_info = next(info for n, info in t2i_ent.items() if info.get("short_id") == "L05")

    shots = []
    for s in data["data"]["scenes"]:
        ve = s.get("visible_entities", [])
        if "L05" not in ve: continue
        si, shi = s["scene_index"], s.get("_shot_index", 1)
        t2i_list = s.get("t2i_variations", [])
        t2i_prompt = t2i_list[0].get("t2i_prompt", "") if t2i_list else ""
        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": list(dict.fromkeys(re.findall(r'C\d{2,3}O\d{2,3}', t2i_prompt))),
            "t2i_prompt": t2i_prompt,
        })

    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


# ──────────────────────────────────────────────
# VLM: blueprint 추출 / 검증
# ──────────────────────────────────────────────

def vlm_extract_blueprint(image_bytes: bytes, set_id: str) -> str:
    """VLM으로 이미지의 공간 구조를 분석해서 blueprint 텍스트 추출."""
    b64 = base64.b64encode(image_bytes).decode("ascii")
    user_prompt = [
        {"type": "text", "text": (
            f"이 이미지는 '{set_id}'라는 세트 배경입니다.\n"
            "이 이미지에 보이는 모든 주요 요소의 **상대적 위치**를 구조적으로 설명하세요.\n"
            "포맷: 각 요소를 '요소명: 프레임 내 위치 (왼쪽/오른쪽/중앙/전경/배경/상단/하단)' 형태로.\n"
            "문, 창문, 싱크대, 냉장고, TV, 침대, 식탁, 가스레인지, 선반 등 고정 설치물 위주.\n"
            "한국어로 답변."
        )},
        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}},
    ]
    result = call_text(
        step="scene_director",  # gemini-pro
        system_prompt="당신은 영화 프로덕션 디자이너입니다. 세트의 공간 구조를 정밀하게 분석합니다.",
        user_prompt=user_prompt,
        temperature=0.1,
    )
    return result


def vlm_verify_consistency(new_image_bytes: bytes, set_id: str, blueprint: str) -> dict:
    """VLM으로 새 이미지가 blueprint와 일치하는지 검증."""
    b64 = base64.b64encode(new_image_bytes).decode("ascii")
    user_prompt = [
        {"type": "text", "text": (
            f"새로 생성된 세트 '{set_id}'가 기존 방 구조(blueprint)와 일치하는지 검증하세요.\n\n"
            f"## 기존 방 구조 (Blueprint)\n{blueprint}\n\n"
            "## 검증 항목\n"
            "1. 문 위치가 같은가?\n"
            "2. 창문 위치/크기가 같은가?\n"
            "3. 주요 가구(싱크대, 냉장고, TV 등) 배치가 같은 벽에 있는가?\n"
            "4. 바닥재/벽지가 같은가?\n\n"
            "JSON으로 답변:\n"
            '{"consistent": true/false, "issues": ["불일치 항목1", ...], "new_elements": ["이 앵글에서 추가로 보이는 요소"]}'
        )},
        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}},
    ]
    result = call_text(
        step="scene_director",
        system_prompt="세트 일관성 검증 전문가. JSON으로만 답변하세요.",
        user_prompt=user_prompt,
        temperature=0.1,
    )
    text = result if isinstance(result, str) else ""
    if "```json" in text: text = text.split("```json")[1].split("```")[0]
    elif "```" in text: text = text.split("```")[1].split("```")[0]
    try:
        return json.loads(text)
    except:
        return {"consistent": True, "issues": [], "new_elements": [], "raw": text}


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

PHASE1_SYSTEM = """You are a professional film production designer.
All shots take place in the SAME physical location.

You will be given the location's EXISTING T2I prompt (entity_t2i).
DO NOT rewrite or modify the location T2I prompt. Use it as-is for parent sets.

Your job:
1. Design BASE SET COMPOSITIONS (max 4, same space max 2 angles).
   For each, specify: camera angle suffix (appended to the existing location T2I).
   Also specify which previous set(s) to use as image reference (max 2, by ID).

2. Identify STATE VARIANTS (children) for each composition.

3. For each shot, SPLIT the T2I prompt into background_prompt + character_prompt.

Rules:
- Max 4 parent compositions total. Same room area = max 2 angles.
- parent t2i_prompt: DO NOT WRITE. The system will construct it as:
  "{entity_t2i} + {camera_angle}" automatically.
  You only provide "camera_angle" — a SHORT camera direction phrase.
  Example: "Wide angle from the bedroom doorway looking inward"
  Example: "Medium angle focusing on the bed and wall area"
- parent ref_ids: max 2 previous parent IDs for visual consistency.
- child t2i_prompt: "Same room as reference image, but with: [state changes only]". NO people. Keep it short.
- character_prompt: ONLY character actions, poses, expressions, held props, camera framing.
  REMOVE all background/room descriptions (already in set reference image).
  NEVER include skin tone/face color modifiers (pale, drained, flushed, ashen, gray face, etc.)
  Express fear/shock through body language only (frozen, trembling, wide eyes, clenched jaw).
  NEVER include gore (blood-soaked, corpse, dead body, exposed flesh, torn) — soften to (motionless figure, slumped, stain, mark).
- child_id = null if room is in default state (use parent directly).
- Output valid JSON. Korean for names, English for prompts."""

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

## Existing Location T2I (entity_t2i — DO NOT MODIFY):
{entity_t2i}

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

## Output JSON:
```json
{{
  "parent_sets": [
    {{
      "id": "SET_A",
      "name": "구도 이름",
      "camera_angle": "Short camera direction phrase only",
      "ref_ids": []
    }}
  ],
  "child_sets": [
    {{
      "id": "SET_A_state",
      "parent_id": "SET_A",
      "state_name": "상태 이름",
      "state_description": "What changed",
      "t2i_prompt": "Same room as reference image, but with: ..."
    }}
  ],
  "shot_assignments": [
    {{
      "shot_label": "S05-Shot1",
      "parent_id": "SET_A",
      "child_id": "SET_A_state or null",
      "character_prompt": "Cleaned T2I — characters/actions ONLY"
    }}
  ]
}}
```"""


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", [])),
        entity_t2i=loc_info.get("t2i_prompt", ""),
        shots_text=shots_text,
    )
    print("[Phase 1] Gemini Pro 분석 중...", flush=True)
    result = call_text(step="scene_director", system_prompt=PHASE1_SYSTEM,
                       user_prompt=user_prompt, temperature=0.1)
    text = result if isinstance(result, str) else ""
    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: 이미지 생성 + VLM 검증
# ──────────────────────────────────────────────

def phase2_generate_parents(parent_sets, entity_t2i: str):
    """Parent 세트 생성: entity_t2i + camera_angle 조합 + VLM 체이닝 검증."""
    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_v4_parent", operation_type="experiment")

    results = {}  # id -> {path, bytes}
    blueprint = ""  # 누적 blueprint

    for p in parent_sets:
        pid = p["id"]
        # entity_t2i 원본 + camera_angle만 추가
        camera_angle = p.get("camera_angle", "")
        prompt = f"{entity_t2i} {camera_angle}".strip()
        ref_ids = p.get("ref_ids", [])[:2]  # 최대 2개

        # 참조 이미지 구성
        labeled_refs = None
        if ref_ids:
            refs_available = [(rid, results[rid]["bytes"]) for rid in ref_ids if rid in results]
            if refs_available:
                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."
                )
                labeled_refs = [(f"Same room ({rid}):", img) for rid, img in refs_available]

        ref_label = f"(refs: {[r for r in ref_ids if r in results]})" if ref_ids else "(no refs)"
        print(f"\n  [parent] {pid} ({p['name']}) {ref_label}")

        # 생성 + 검증 (max 2 retries)
        for attempt in range(3):
            print(f"    생성 (attempt {attempt+1})... ", end="", flush=True)
            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")
            except Exception as e:
                print(f"FAIL: {e}")
                continue

            # VLM 검증
            if not blueprint:
                # 첫 번째: blueprint 추출
                print(f"    VLM blueprint 추출... ", end="", flush=True)
                blueprint = vlm_extract_blueprint(img_bytes, pid)
                print(f"OK ({len(blueprint)}자)")
                print(f"    Blueprint: {blueprint[:200]}...")
                results[pid] = {"path": out, "bytes": img_bytes}
                break
            else:
                # 이후: 일관성 검증
                print(f"    VLM 검증... ", end="", flush=True)
                verify = vlm_verify_consistency(img_bytes, pid, blueprint)
                consistent = verify.get("consistent", True)
                issues = verify.get("issues", [])
                new_elements = verify.get("new_elements", [])

                if consistent or attempt == 2:
                    if issues:
                        print(f"경미한 불일치: {issues}")
                    else:
                        print("OK ✓")
                    # blueprint 보강
                    if new_elements:
                        blueprint += f"\n\n[{pid}에서 추가 확인]: {', '.join(new_elements)}"
                        print(f"    Blueprint 보강: +{len(new_elements)}개 요소")
                    results[pid] = {"path": out, "bytes": img_bytes}
                    break
                else:
                    print(f"불일치! {issues} → 재생성")

    print(f"\n  최종 Blueprint ({len(blueprint)}자):\n  {blueprint[:300]}...")
    return results, blueprint


def phase2_generate_children(child_sets, parent_results, blueprint):
    """Child 세트 생성 (parent 참조 + VLM 검증)."""
    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_v4_child", operation_type="experiment")

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

        prompt = c["t2i_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 exact room structure:", parent["bytes"])]

        print(f"  [child] {cid} ({c['state_name']}) ← {parent_id}... ", end="", flush=True)

        for attempt in range(2):
            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)
                elapsed = time.time() - t0
                print(f"OK {len(img_bytes)//1024}KB {elapsed:.1f}s", end="")

                # VLM 검증 (parent와 구조 동일한지)
                verify = vlm_verify_consistency(img_bytes, cid, blueprint)
                if verify.get("consistent", True) or attempt == 1:
                    issues = verify.get("issues", [])
                    if issues:
                        print(f" (경미: {issues[:2]})")
                    else:
                        print(" ✓")
                    results[cid] = {"path": out, "bytes": img_bytes}
                    break
                else:
                    print(f" 불일치→재생성")
            except Exception as e:
                print(f"FAIL: {e}")
                break

    return results


def get_character_refs(db, ve, t2i_text):
    """character_prompt에 언급된 인물/소품만 참조 로드."""
    labeled_refs = []
    for sid in ve:
        if sid.startswith("L"): continue
        # character_prompt에 언급 안 되면 스킵
        if sid.startswith("C") and not re.search(rf'(?<![CO\d]){sid}', t2i_text):
            continue
        if sid.startswith("P") and not re.search(rf'(?<![A-Z]){sid}(?!\d)', t2i_text):
            continue

        entity = db.query(Entity).filter(Entity.project_id == PROJECT_ID, Entity.short_id == sid).first()
        if not entity: continue
        eid, etype = entity.id, entity.entity_type

        if etype == "character":
            m = re.search(rf'{sid}(O\d{{2,3}})', t2i_text)
            if m:
                outlook = db.query(Entity).filter(Entity.project_id == PROJECT_ID, Entity.short_id == m.group(1)).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}{m.group(1)} — 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_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 sanitize_gore(text):
    """SAFETY 블록 방지를 위한 gore 표현 완화."""
    text = re.sub(r'blood-soaked torn shoulder and collarbone of the slumped corpse',
                  'a motionless figure slumped behind the curtain', text)
    text = re.sub(r'dead woman.*?running over exposed flesh',
                  'motionless figure sitting on the floor with head hanging forward', text)
    text = re.sub(r'blood-soaked|corpse|dead body|dead woman|exposed flesh|torn shoulder',
                  '', text, flags=re.IGNORECASE)
    text = re.sub(r'face drained of color|pale face|ashen face|gray face|white face|flushed face',
                  'shocked expression', text, flags=re.IGNORECASE)
    text = re.sub(r'\s+', ' ', text).strip()
    return text


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_v4_shot", operation_type="experiment",
                       project_id=PROJECT_ID, episode_id=EPISODE_ID)

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

            # 배경 선택: 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")
                continue

            shot = next((s for s in shots_data if s["label"] == label), None)
            if not shot: continue

            entity_refs = get_character_refs(db, shot["visible_entities"], char_prompt)
            labeled_refs = [("SET BACKGROUND — use this exact room as the background", bg["bytes"])]
            labeled_refs.extend(entity_refs)

            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)
                (shot_dir / f"{label.replace('-','_')}.png").write_bytes(img_bytes)
                print(f"OK {len(img_bytes)//1024}KB {time.time()-t0:.1f}s")
            except Exception as e:
                # SAFETY fallback: child → parent
                if child_id and child_id in child_results and parent_id in parent_results:
                    print(f"BLOCKED → parent fallback... ", end="", flush=True)
                    labeled_refs[0] = ("SET BACKGROUND — use this exact room as the background", parent_results[parent_id]["bytes"])
                    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,
                    )
                    try:
                        img_bytes, _ = client.generate_image(
                            prompt=final_prompt, aspect_ratio="16:9", labeled_references=labeled_refs)
                        (shot_dir / f"{label.replace('-','_')}.png").write_bytes(img_bytes)
                        print(f"OK {len(img_bytes)//1024}KB {time.time()-t0:.1f}s (parent)")
                    except Exception as e2:
                        print(f"FAIL: {e2}")
                else:
                    print(f"FAIL: {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"{'='*60}")
    print(f"Set Design v4 — L05 ({loc_info['name']}) | {len(shots)} shots")
    print(f"{'='*60}\n")

    # Phase 1
    print("[Phase 1] LLM 분석")
    analysis = phase1_analyze(loc_info, shots, names)
    parents = analysis["parent_sets"]
    children = analysis["child_sets"]
    assignments = analysis["shot_assignments"]

    print(f"\n  Parents: {len(parents)}개 — {[p['id'] for p in parents]}")
    print(f"  Children: {len(children)}개 — {[c['id'] for c in children]}")
    for a in assignments:
        bg = a.get("child_id") or a["parent_id"]
        print(f"    {a['shot_label']} → {bg}")

    with open(OUT_DIR / "phase1_analysis.json", "w") as f:
        json.dump(analysis, f, ensure_ascii=False, indent=2)

    # Phase 2
    print(f"\n[Phase 2-1] Parent 생성 + VLM 체이닝 검증")
    parent_results, blueprint = phase2_generate_parents(parents, loc_info.get("t2i_prompt", ""))

    with open(OUT_DIR / "blueprint.txt", "w") as f:
        f.write(blueprint)

    print(f"\n[Phase 2-2] Child 생성 + VLM 검증")
    child_results = phase2_generate_children(children, parent_results, blueprint)

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

    # URLs
    base_url = "http://192.168.231.91:3002/experiment/set_v4"
    print(f"\n{'='*60}")
    print("URLs")
    print(f"{'='*60}")
    print("\nParent:")
    for p in sorted((OUT_DIR / "parents").glob("*.png")):
        print(f"  {base_url}/parents/{p.name}")
    print("\nChildren:")
    for p in sorted((OUT_DIR / "children").glob("*.png")):
        print(f"  {base_url}/children/{p.name}")
    print("\nShots:")
    for p in sorted((OUT_DIR / "shots").glob("*.png")):
        print(f"  {base_url}/shots/{p.name}")


if __name__ == "__main__":
    main()