"""실험 v9: 6-Phase 세트 디자인 파이프라인 (와이드 1장 base)

Phase 1: 전체 T2I → LLM → 와이드 1장 기초 배경 T2I → 이미지 생성
Phase 2: 샷 T2I + 기초 T2I → LLM → 최대 4개 배경 변형 T2I → 이미지 생성
Phase 3: 샷 T2I + 변형 T2I → LLM → 각 샷: 앞쪽 샷 참조 or 배경 이미지 선택
Phase 4: 앞쪽 참조 시 → 두 T2I 비교 → LLM → 삭제 대상 판별 (개별 호출)
Phase 5+6: 삭제 실행 + 최종 샷 생성 (순차)

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

import json
import pathlib
import re
import sys
import time
import urllib.request

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.llm_client import call_text
from app.modules.llm.gemini_image_client import GeminiImageClient
from app.services.prompt_service import build_final_scene_prompt as _build_final_scene_prompt

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


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

def load_shots():
    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)

    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

    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,
            "visible_entities": ve,
            "t2i_prompt": t2i_prompt,
            "beat_title": s.get("beat_title", ""),
            "chars": [e for e in ve if e.startswith("C")],
            "props": [e for e in ve if e.startswith("P")],
        })
    return shots, names, etm


# ──────────────────────────────────────────────
# LLM 호출 (Gemini Pro → GPT fallback)
# ──────────────────────────────────────────────

def _call_llm(prompt, temperature=0.1, max_retry=3):
    """Gemini Pro 직접 호출 → 실패 시 GPT-5.4 fallback."""
    api_key = settings.gemini_api_key
    model = "gemini-3.1-pro-preview"
    url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}"

    body = {
        "contents": [{"role": "user", "parts": [{"text": prompt}]}],
        "generationConfig": {"temperature": temperature},
        "safetySettings": [
            {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
            {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
            {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
            {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
        ],
    }

    for attempt in range(max_retry):
        try:
            req = urllib.request.Request(
                url, data=json.dumps(body).encode(),
                headers={"Content-Type": "application/json"}, method="POST",
            )
            resp = urllib.request.urlopen(req, timeout=120)
            payload = json.loads(resp.read())
            if "candidates" in payload:
                text = payload["candidates"][0]["content"]["parts"][0].get("text", "")
                if text.strip():
                    return text
            block = payload.get("promptFeedback", {}).get("blockReason", "")
            print(f"  (Gemini 차단: {block}, {attempt+1}/{max_retry})", flush=True)
        except Exception as e:
            print(f"  (Gemini 에러: {e}, {attempt+1}/{max_retry})", flush=True)
        time.sleep(3)

    print("  → GPT-5.4 fallback", flush=True)
    result = call_text(step="entity_detail_batch", system_prompt="", user_prompt=prompt, temperature=temperature)
    return result if isinstance(result, str) else ""


def _parse_json(text, debug_name=""):
    if debug_name:
        (OUT_DIR / f"{debug_name}_raw.txt").write_text(text, encoding="utf-8")
    if "```json" in text:
        text = text.split("```json")[1].split("```")[0]
    elif "```" in text:
        text = text.split("```")[1].split("```")[0]
    text = text.strip()
    if not text:
        raise ValueError(f"Empty JSON after parsing. Raw length: {len(text)}")
    return json.loads(text)


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


def _generate_with_retry(client, prompt, labeled_refs, max_retry=3):
    for attempt in range(max_retry):
        try:
            img_bytes, _ = client.generate_image(
                prompt=prompt, aspect_ratio="16:9", labeled_references=labeled_refs,
            )
            return img_bytes
        except Exception as e:
            is_safety = "SAFETY" in str(e) or "moderation" in str(e)
            if attempt < max_retry - 1 and is_safety:
                time.sleep(3)
                continue
            if attempt == max_retry - 1:
                if not is_safety:
                    print(f"  ERROR (non-safety): {e}")
                return None
            raise


# ──────────────────────────────────────────────
# Phase 1: 와이드 1장 기초 배경
# ──────────────────────────────────────────────

def phase1_base_image(shots, client):
    """전체 샷 T2I → LLM → 와이드 1장 T2I → 이미지 생성."""

    shots_block = "\n".join(
        f"[{s['label']}]: {s['t2i_prompt']}" for s in shots
    )

    prompt = f"""You are a film set designer analyzing a fictional Korean film screenplay storyboard (영화 시나리오 콘티 분석).
The following are text-to-image generation prompts from a fictional thriller movie script. These describe planned storyboard frames — no real events.

Design a SINGLE SPLIT IMAGE with TWO SIDE-BY-SIDE PANELS showing the same location.
- LEFT PANEL: one half/area of the location (normal camera angle, not panoramic)
- RIGHT PANEL: the other half/area (normal camera angle, not panoramic)
- Extract every PURE BACKGROUND element from ALL shots: furniture, curtains, TV, lamps, table, chairs, bed, desk, door, sink, cabinets, mirror, heater, partition curtain, etc.
- Both panels must share identical wall color, floor texture, ceiling
- NO characters, NO hand-held props
- EXCLUDE story-event changes (blood, mess, overturned items)
- Include "Clear vertical dividing line between the two panels" in the T2I

## Fictional film storyboard T2I prompts:
{shots_block}

Output JSON: {{"location_understanding": "...", "left_panel": "description of left panel area", "right_panel": "description of right panel area", "base_t2i": "A single image divided into TWO SIDE-BY-SIDE PANELS showing the same ... LEFT PANEL: ... RIGHT PANEL: ... Both panels share identical ... Clear vertical dividing line ... Photorealistic cinematic still."}}"""

    print("\n[Phase 1] LLM: 분할 기초 배경 T2I 생성...", flush=True)
    raw = _call_llm(prompt)
    result = _parse_json(raw, "phase1")
    _save("phase1.json", result)

    t2i = result["base_t2i"]
    print(f"  T2I: {t2i[:100]}...", flush=True)
    print(f"  생성 중... ", end="", flush=True)
    t0 = time.time()
    img_bytes = _generate_with_retry(client, t2i, [])
    if img_bytes:
        (OUT_DIR / "base" / "split.png").write_bytes(img_bytes)
        print(f"OK {len(img_bytes)//1024}KB {time.time()-t0:.1f}s")
    else:
        print("FAIL")
    return result, img_bytes


# ──────────────────────────────────────────────
# Phase 2: 배경 변형 최대 4장
# ──────────────────────────────────────────────

def phase2_bg_variants(shots, phase1_result, base_bytes, client):
    """샷 T2I + 기초 T2I → 최대 4개 배경 변형 생성."""

    base_t2i = phase1_result["base_t2i"]

    shots_block = "\n".join(
        f"[{s['label']}]: {s['t2i_prompt']}"
        for s in shots
    )

    prompt = f"""You are a film set designer working on a fictional film screenplay storyboard (영화 시나리오 콘티 분석).

You have 1 base wide-angle background image (BASE) of the full apartment:
BASE: {base_t2i}...

Analyze the shot T2I prompts and identify background STATE CHANGES that need variant images:
- Stains/marks appearing, room wrecked, items overturned
- Lighting changes (day→night)
- Curtain opened/closed
- Room cleaned after being dirty

Rules:
- At most 4 variants (BG_1 to BG_4)
- Each variant references BASE
- T2I: ONLY background changes, no characters, no props
- Include which AREA of the room is affected (kitchen side, bedroom side, or full room)

## Fictional film storyboard T2I prompts:
{shots_block}

Output JSON: {{"variants": [{{"variant_id": "BG_1", "name": "short name", "area": "bedroom/kitchen/full", "t2i_prompt": "Photorealistic cinematic still. ...", "change_summary": "what changed"}}]}}"""

    print("\n[Phase 2] LLM: 배경 변형 T2I 생성...", flush=True)
    raw = _call_llm(prompt)
    result = _parse_json(raw, "phase2")
    _save("phase2.json", result)

    generated = {}
    for v in result.get("variants", []):
        vid = v["variant_id"]
        t2i = v["t2i_prompt"]
        refs = []
        if base_bytes:
            refs.append(("BASE — full apartment, maintain architecture, apply changes", base_bytes))

        print(f"  [{vid}] {v['name']} 생성... ", end="", flush=True)
        t0 = time.time()
        img_bytes = _generate_with_retry(client, t2i, refs)
        if img_bytes:
            (OUT_DIR / "bg" / f"{vid.lower()}.png").write_bytes(img_bytes)
            generated[vid] = img_bytes
            print(f"OK {len(img_bytes)//1024}KB {time.time()-t0:.1f}s")
        else:
            print("FAIL")

    return result, generated


# ──────────────────────────────────────────────
# Phase 3: 샷별 배경 선택
# ──────────────────────────────────────────────

def phase3_shot_assignments(shots, phase1_result, phase2_result):
    """각 샷: 앞쪽 샷 참조 or 배경 이미지(BASE/variant) 선택."""

    bg_list = [f"BASE: full apartment panorama — {phase1_result['base_t2i']}"]
    for v in phase2_result.get("variants", []):
        bg_list.append(f"{v['variant_id']}: {v['name']} — {v['change_summary']}")
    bg_block = "\n".join(bg_list)

    shots_block = "\n".join(
        f"[{i}] {s['label']}: {s['t2i_prompt']}"
        for i, s in enumerate(shots)
    )

    prompt = f"""You are a film director working on a fictional film screenplay storyboard (영화 시나리오 콘티 분석).

For each shot, choose ONE background reference:
- **"background"**: Use BASE or a variant (BG_1, BG_2, etc.)
- **"prev_shot"**: Use a previous shot's generated image (background carries forward, only characters changed)

Use "prev_shot" only when background is essentially identical to a previous shot.
Use "background" when the background state matches a base/variant image.
IMPORTANT: The first shot (index 0) MUST use "background". A "prev_shot" ref_shot must be an EARLIER shot in this list, never a later one.

## Available backgrounds:
{bg_block}

## Shots (in story order):
{shots_block}

Output JSON: {{"assignments": [{{"shot_label": "S05_Shot1", "type": "background", "bg_id": "BASE", "reason": "brief"}}, {{"shot_label": "S12_Shot2", "type": "prev_shot", "ref_shot": "S12_Shot1", "reason": "brief"}}]}}"""

    print("\n[Phase 3] LLM: 샷별 배경 선택...", flush=True)
    raw = _call_llm(prompt)
    result = _parse_json(raw, "phase3")

    # forward reference 검증: prev_shot은 반드시 앞쪽 샷만
    shot_labels = [s["label"] for s in shots]
    label_order = {label: i for i, label in enumerate(shot_labels)}
    for a in result.get("assignments", []):
        if a["type"] == "prev_shot":
            cur_idx = label_order.get(a["shot_label"], -1)
            ref_idx = label_order.get(a.get("ref_shot", ""), -1)
            if ref_idx < 0 or ref_idx >= cur_idx:
                print(f"  WARNING: {a['shot_label']} → prev:{a.get('ref_shot')} forward ref → BASE로 교체")
                a["type"] = "background"
                a["bg_id"] = "BASE"

    _save("phase3.json", result)
    return result


# ──────────────────────────────────────────────
# Phase 4: 앞쪽 샷 참조 시 삭제 대상 판별
# ──────────────────────────────────────────────

def phase4_removal_analysis(shots, phase3_result):
    """prev_shot 참조: 두 T2I 비교 → 삭제 대상 판별 (개별 호출)."""
    shot_map = {s["label"]: s for s in shots}
    assignments = phase3_result.get("assignments", [])

    prev_refs = [a for a in assignments if a["type"] == "prev_shot"]
    if not prev_refs:
        print("\n[Phase 4] prev_shot 참조 없음 — 스킵")
        _save("phase4.json", {"removals": []})
        return {"removals": []}

    system = """You are a film continuity supervisor working on a fictional film screenplay storyboard (영화 시나리오 콘티 분석).
Compare two consecutive shot T2I prompts. Identify elements in the PREVIOUS shot but NOT in the CURRENT shot.
Focus on: characters/people who left, objects that disappeared.
Output JSON only."""

    removals = []
    print(f"\n[Phase 4] 삭제 대상 판별 ({len(prev_refs)}건)...", flush=True)

    for a in prev_refs:
        cur = shot_map.get(a["shot_label"])
        ref = shot_map.get(a["ref_shot"])
        if not cur or not ref:
            continue

        user = f"""(Fictional film storyboard analysis.)

## Previous shot ({a['ref_shot']}):
T2I: {ref['t2i_prompt']}

## Current shot ({a['shot_label']}):
T2I: {cur['t2i_prompt']}

Output JSON: {{"shot_label": "{a['shot_label']}", "ref_shot": "{a['ref_shot']}", "removal_needed": true/false, "removal_prompt": "Remove [desc] from the image. Fill with natural background.", "removed_elements": ["elem1"]}}"""

        print(f"  {a['shot_label']} ← {a['ref_shot']}... ", end="", flush=True)
        raw = _call_llm(f"{system}\n\n{user}")
        try:
            r = _parse_json(raw)
            removals.append(r)
            if r.get("removal_needed"):
                print(f"삭제: {r.get('removed_elements', [])}")
            else:
                print("변경 없음")
        except Exception as e:
            print(f"FAIL: {e}")

    result = {"removals": removals}
    _save("phase4.json", result)
    return result


# ──────────────────────────────────────────────
# 유틸리티
# ──────────────────────────────────────────────

def get_composite_refs(db, ve, t2i_text):
    """composite 우선 참조 이미지 로드."""
    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|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',
                  'shocked expression', text, flags=re.IGNORECASE)
    return re.sub(r'\s+', ' ', text).strip()


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

def main():
    for d in ["base", "bg", "cleaned", "shots"]:
        (OUT_DIR / d).mkdir(parents=True, exist_ok=True)

    shots, names, etm = load_shots()
    print(f"{'='*60}")
    print(f"Set Design v9 — 6-Phase Pipeline (Wide Base)")
    print(f"L05 | {len(shots)} shots")
    print(f"{'='*60}")
    for s in shots:
        print(f"  {s['label']} | {s['beat_title'][:40]}")

    client = GeminiImageClient(model=settings.gemini_image_model)
    client.set_context(step="set_v9", operation_type="experiment",
                       project_id=PROJECT_ID, episode_id=EPISODE_ID)
    db = SessionLocal()

    try:
        # Phase 1: 와이드 1장
        p1_result, base_bytes = phase1_base_image(shots, client)

        # Phase 2: 변형 최대 4장
        p2_result, bg_variants = phase2_bg_variants(shots, p1_result, base_bytes, client)

        # Phase 3: 샷별 배정
        p3_result = phase3_shot_assignments(shots, p1_result, p2_result)

        assignments = p3_result.get("assignments", [])
        bg_count = sum(1 for a in assignments if a["type"] == "background")
        prev_count = sum(1 for a in assignments if a["type"] == "prev_shot")
        print(f"\n  배정: background={bg_count}, prev_shot={prev_count}")
        for a in assignments:
            if a["type"] == "background":
                print(f"    {a['shot_label']} → {a['bg_id']}")
            else:
                print(f"    {a['shot_label']} → prev:{a['ref_shot']}")

        # Phase 4: 삭제 판별
        p4_result = phase4_removal_analysis(shots, p3_result)

        # Phase 5+6: 순차 생성
        all_bg = {"BASE": base_bytes, **bg_variants}
        assign_map = {a["shot_label"]: a for a in assignments}
        removal_map = {r["shot_label"]: r for r in p4_result.get("removals", []) if r.get("removal_needed")}
        shot_images = {}
        cleaned_images = {}

        print(f"\n{'='*60}")
        print("Phase 5+6: 샷 이미지 순차 생성")
        print(f"{'='*60}")

        for shot in shots:
            label = shot["label"]
            assign = assign_map.get(label)
            if not assign:
                continue

            if assign["type"] == "prev_shot":
                ref_label = assign["ref_shot"]
                ref_bytes = shot_images.get(ref_label)

                # Phase 5: 삭제
                if label in removal_map and ref_bytes:
                    r = removal_map[label]
                    refs = [("Source image — remove specified elements:", ref_bytes)]
                    prompt = f"{r.get('removal_prompt', '')} Keep everything else exactly the same. Photorealistic."

                    print(f"  [{label}] 정제 ({ref_label})... ", end="", flush=True)
                    t0 = time.time()
                    cleaned = _generate_with_retry(client, prompt, refs)
                    if cleaned:
                        (OUT_DIR / "cleaned" / f"{label}_cleaned.png").write_bytes(cleaned)
                        cleaned_images[label] = cleaned
                        print(f"OK {time.time()-t0:.1f}s")
                    else:
                        print("FAIL")

                bg_bytes = cleaned_images.get(label) or ref_bytes
                bg_label = f"cleaned:{ref_label}" if label in cleaned_images else ref_label
            else:
                bg_id = assign.get("bg_id", "BASE")
                bg_bytes = all_bg.get(bg_id)
                bg_label = bg_id

            if not bg_bytes:
                print(f"  [{label}] 배경 없음 — 스킵")
                continue

            # Phase 6: 최종 샷 생성
            char_refs = get_composite_refs(db, shot["visible_entities"], shot["t2i_prompt"])
            labeled_refs = [(f"SET BACKGROUND ({bg_label})", bg_bytes)]
            labeled_refs.extend(char_refs)

            final_prompt = _build_final_scene_prompt(
                t2i_prompt=sanitize_gore(shot["t2i_prompt"]),
                labeled_refs=labeled_refs,
                style_context="",
                scene_index=shot["scene_index"],
                entity_text_map=etm,
            )

            print(f"  [{label}] bg={bg_label} refs=1+{len(char_refs)}... ", end="", flush=True)
            t0 = time.time()
            img_bytes = _generate_with_retry(client, final_prompt, labeled_refs)
            if img_bytes:
                (OUT_DIR / "shots" / f"{label}.png").write_bytes(img_bytes)
                shot_images[label] = img_bytes
                print(f"OK {len(img_bytes)//1024}KB {time.time()-t0:.1f}s")
            else:
                print("FAIL")

    finally:
        db.close()

    # URLs
    base_url = "http://192.168.231.91:3000/experiment/set_v9"
    print(f"\n{'='*60}")
    print("URLs")
    print(f"{'='*60}")
    print("\nBase:")
    for p in sorted((OUT_DIR / "base").glob("*.png")):
        print(f"  {base_url}/base/{p.name}")
    print("\nBG Variants:")
    for p in sorted((OUT_DIR / "bg").glob("*.png")):
        print(f"  {base_url}/bg/{p.name}")
    print("\nCleaned:")
    for p in sorted((OUT_DIR / "cleaned").glob("*.png")):
        print(f"  {base_url}/cleaned/{p.name}")
    print("\nShots:")
    for p in sorted((OUT_DIR / "shots").glob("*.png")):
        print(f"  {base_url}/shots/{p.name}")


if __name__ == "__main__":
    main()