"""실험 v5: 4-View Grid 피라미드 세트 디자인

피라미드 구조:
  [1] 4-View Grid (1장) — 방 전체 구조의 유일한 진실
  [2] Parent sets — grid 참조 → 개별 앵글 (grid 전체 이미지를 참조로 전달)
  [3] Child sets — parent 참조 → 상태 변형
  [4] Shot images — child/parent + 인물 refs → 최종

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

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

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


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

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


# ──────────────────────────────────────────────
# Layer 1: 4-View Grid 생성
# ──────────────────────────────────────────────

def generate_grid(loc_info: dict, shots: list) -> bytes:
    """전체 샷의 배경 요소를 종합해서 4-view grid 1장 생성."""
    client = GeminiImageClient(model=settings.gemini_image_model)
    client.set_context(step="set_v5_grid", operation_type="experiment")

    # 실내/야외 판단
    desc = loc_info.get("description", "")
    is_indoor = any(kw in desc for kw in ["내부", "실내", "방", "사무실", "조타실"])

    entity_t2i = loc_info.get("t2i_prompt", "")
    # "no people, no vehicles" 등 제거 — grid 프롬프트에서 별도 지시
    entity_t2i_clean = re.sub(r',?\s*no people.*$', '', entity_t2i, flags=re.IGNORECASE).strip().rstrip(',.')

    if is_indoor:
        # 샷에서 배경 요소 추출 (LLM 없이 간단히)
        all_bg_keywords = set()
        for s in shots:
            t2i = s["t2i_prompt"]
            # [L05: ...] 안의 내용 추출
            for m in re.finditer(r'\[L\d+:\s*([^\]]+)\]', t2i):
                words = m.group(1).split(',')
                for w in words:
                    w = w.strip().lower()
                    if any(kw in w for kw in ['room', 'wall', 'floor', 'window', 'curtain', 'bed', 'door',
                                               'kitchen', 'sink', 'tv', 'table', 'lamp', 'desk', 'cabinet',
                                               'stove', 'shelf', 'refrigerator', 'ceiling']):
                        all_bg_keywords.add(w.strip())

        prompt = (
            f"{entity_t2i_clean}.\n\n"
            "Generate a SINGLE image divided into 4 EQUAL QUADRANTS (2x2 grid).\n"
            "Each quadrant shows the SAME room from the CENTER, looking in a different direction:\n\n"
            "TOP-LEFT: Looking toward the kitchen/sink wall\n"
            "TOP-RIGHT: Looking toward the main window wall (with TV and table area)\n"
            "BOTTOM-LEFT: Looking toward the entrance door\n"
            "BOTTOM-RIGHT: Looking toward the bedroom area (through doorway if separate, or bed area)\n\n"
            "All 4 views must show the SAME room — identical wall color, flooring, ceiling.\n"
            "Photorealistic cinematic interior. No people, no figures, empty room only."
        )
    else:
        prompt = (
            f"{entity_t2i_clean}.\n\n"
            "Generate a SINGLE image divided into 4 EQUAL QUADRANTS (2x2 grid).\n"
            "Each quadrant shows the SAME location from DIFFERENT camera positions looking inward:\n\n"
            "TOP-LEFT: From the north looking south\n"
            "TOP-RIGHT: From the east looking west\n"
            "BOTTOM-LEFT: From the south looking north\n"
            "BOTTOM-RIGHT: From the west looking east\n\n"
            "All 4 views must show the SAME location — consistent weather, time, structures.\n"
            "Photorealistic cinematic. No people, no figures."
        )

    grid_dir = OUT_DIR / "grid"
    grid_dir.mkdir(parents=True, exist_ok=True)

    print("[Layer 1] 4-View Grid 생성... ", end="", flush=True)
    t0 = time.time()
    img_bytes, _ = client.generate_image(prompt=prompt, aspect_ratio="1:1")
    (grid_dir / "grid.png").write_bytes(img_bytes)
    print(f"OK {len(img_bytes)//1024}KB {time.time()-t0:.1f}s")
    return img_bytes


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

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

A 4-VIEW GRID image of this location already exists (2x2: TL=kitchen, TR=window/TV, BL=entrance, BR=bedroom).
You must assign each shot to one of the 4 quadrants and design state variants.

Your job:
1. Assign each quadrant a SET ID (SET_TL, SET_TR, SET_BL, SET_BR).
   Give each a short Korean name.

2. Identify STATE VARIANTS (children) needed.
   States = story-driven changes (blood, mess, unnaturally clean, etc.)
   Each child references a specific quadrant.

3. For each shot, SPLIT the T2I prompt into:
   - character_prompt: ONLY character actions, poses, expressions, held props, camera framing.
     REMOVE all background/room descriptions.
     NEVER include skin tone/face color modifiers (pale, drained, flushed, ashen, gray).
     Express fear/shock through body language only.
     NEVER include gore — soften to (motionless figure, slumped, stain, mark).
   - Assign parent quadrant + child (or null if default state).

Rules:
- child t2i_prompt: "Same room as reference image, but with: [changes]". NO people. Keep short.
- child_id = null if room is in default state.
- Output valid JSON. Korean for names, English for prompts."""

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

## 4-View Grid quadrants:
- TOP-LEFT (SET_TL): Kitchen/sink wall
- TOP-RIGHT (SET_TR): Window/TV/table wall
- BOTTOM-LEFT (SET_BL): Entrance door
- BOTTOM-RIGHT (SET_BR): Bedroom area

## All shots:
{shots_text}

## Output JSON:
```json
{{
  "quadrants": [
    {{"id": "SET_TL", "name": "주방 방향"}},
    {{"id": "SET_TR", "name": "창문/TV 방향"}},
    {{"id": "SET_BL", "name": "현관 방향"}},
    {{"id": "SET_BR", "name": "침실 방향"}}
  ],
  "child_sets": [
    {{
      "id": "SET_BR_bloody",
      "parent_id": "SET_BR",
      "state_name": "혈흔 상태",
      "t2i_prompt": "Same room as reference image, but with: ..."
    }}
  ],
  "shot_assignments": [
    {{
      "shot_label": "S05-Shot1",
      "parent_id": "SET_TL",
      "child_id": null,
      "character_prompt": "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']} | {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"],
        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)


# ──────────────────────────────────────────────
# Layer 2: Parent 생성 (grid 참조)
# ──────────────────────────────────────────────

QUADRANT_DIRECTIONS = {
    "SET_TL": "the TOP-LEFT quadrant (kitchen/sink direction)",
    "SET_TR": "the TOP-RIGHT quadrant (window/TV direction)",
    "SET_BL": "the BOTTOM-LEFT quadrant (entrance door direction)",
    "SET_BR": "the BOTTOM-RIGHT quadrant (bedroom direction)",
}


def generate_parents(quadrants, grid_bytes, loc_info):
    """Grid 참조로 개별 앵글 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_v5_parent", operation_type="experiment")

    entity_t2i = loc_info.get("t2i_prompt", "")
    entity_t2i_clean = re.sub(r',?\s*no people.*$', '', entity_t2i, flags=re.IGNORECASE).strip().rstrip(',.')

    # 사용되는 quadrant만 생성
    used_ids = set(q["id"] for q in quadrants)
    results = {}

    for q in quadrants:
        qid = q["id"]
        direction = QUADRANT_DIRECTIONS.get(qid, "")

        prompt = (
            f"{entity_t2i_clean}. "
            f"Use {direction} of the reference grid image as the camera angle and composition. "
            f"Generate a full 16:9 cinematic shot matching that quadrant's view of the room. "
            f"Maintain the exact same room structure, wall colors, flooring, furniture as shown in the grid. "
            f"No people, empty room only."
        )

        labeled_refs = [("4-view grid of this room — match the specified quadrant:", grid_bytes)]

        print(f"  [parent] {qid} ({q['name']})... ", 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"{qid}.png"
            out.write_bytes(img_bytes)
            print(f"OK {len(img_bytes)//1024}KB {time.time()-t0:.1f}s")
            results[qid] = {"path": out, "bytes": img_bytes}
        except Exception as e:
            print(f"FAIL: {e}")

    return results


# ──────────────────────────────────────────────
# Layer 3: Child 생성 (parent 참조)
# ──────────────────────────────────────────────

def generate_children(child_sets, parent_results):
    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_v5_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 = (
            f"{c['t2i_prompt']} "
            "Keep the exact same room structure, walls, floor, furniture layout as the reference image. "
            "Only apply the described changes. No people."
        )
        labeled_refs = [("Base room — keep this exact 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


# ──────────────────────────────────────────────
# Layer 4: 샷 이미지 생성
# ──────────────────────────────────────────────

def get_character_refs(db, ve, t2i_text):
    labeled_refs = []
    for sid in ve:
        if sid.startswith("L"): continue
        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):
    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)
    return re.sub(r'\s+', ' ', text).strip()


def 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_v5_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"])

            if child_id and child_id in child_results:
                bg, bg_label = child_results[child_id], child_id
            elif parent_id in parent_results:
                bg, bg_label = parent_results[parent_id], 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:
                if child_id and parent_id in parent_results:
                    print(f"BLOCKED → parent fallback... ", end="", flush=True)
                    labeled_refs[0] = ("SET 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 (parent) {len(img_bytes)//1024}KB")
                    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 v5 (Grid Pyramid) — L05 ({loc_info['name']}) | {len(shots)} shots")
    print(f"{'='*60}\n")

    # Layer 1: Grid
    grid_bytes = generate_grid(loc_info, shots)

    # Phase 1: LLM Analysis
    print(f"\n[Phase 1] LLM 분석")
    analysis = phase1_analyze(loc_info, shots, names)
    quadrants = analysis["quadrants"]
    children = analysis["child_sets"]
    assignments = analysis["shot_assignments"]

    print(f"\n  Quadrants: {[q['id'] + '=' + q['name'] for q in quadrants]}")
    print(f"  Children: {len(children)}개")
    for c in children:
        print(f"    {c['id']} ← {c['parent_id']}: {c['state_name']}")
    print(f"  Assignments:")
    for a in assignments:
        bg = a.get("child_id") or a["parent_id"]
        print(f"    {a['shot_label']} → {bg}")

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

    # Layer 2: Parents (grid 참조)
    print(f"\n[Layer 2] Parent 생성 (grid 참조)")
    # 실제 사용되는 quadrant만
    used_parents = set(a["parent_id"] for a in assignments)
    used_parents.update(c["parent_id"] for c in children)
    active_quadrants = [q for q in quadrants if q["id"] in used_parents]
    parent_results = generate_parents(active_quadrants, grid_bytes, loc_info)

    # Layer 3: Children (parent 참조)
    print(f"\n[Layer 3] Child 생성 (parent 참조)")
    child_results = generate_children(children, parent_results)

    # Layer 4: Shots
    print(f"\n[Layer 4] 샷 생성")
    generate_shots(assignments, parent_results, child_results, shots, etm)

    # URLs
    base_url = "http://192.168.231.91:3002/experiment/set_v5"
    print(f"\n{'='*60}")
    print("URLs")
    print(f"{'='*60}")
    print(f"\nGrid: {base_url}/grid/grid.png")
    print("\nParents:")
    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()