"""실험 v8: 체이닝 배경 + 점선 실루엣 + LLM 판단 기반 요소 제거

구조:
1. 기본 세트 2장 생성 (v7과 동일)
2. 매 샷마다 LLM이 판단:
   - 배경 변화 있는가? (이전 샷 T2I vs 현재 샷 T2I)
   - 변화 있으면: 이전 샷 이미지 참조
   - 사라지는 요소 있으면: 이전 샷 이미지에서 요소 제거 → 정제 배경
3. 배경 참조에 점선 실루엣 추가 (현재 샷 인물 수/위치 기반)
4. 최종 샷: 배경 참조 + composite refs

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

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


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

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

    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", ""),
            "character_combos": list(dict.fromkeys(re.findall(r'C\d{2,3}O\d{2,3}', t2i_prompt))),
            "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: 전체 샷 시퀀스 분석
# ──────────────────────────────────────────────

def analyze_shot_sequence(shots, names):
    """LLM이 전체 샷 시퀀스를 보고 매 샷의 배경 전략을 결정."""

    system = """You are a film director planning background image generation for a sequence of shots in the same location.

For each shot, you must decide:
1. background_source: "base_set" (use the base set image) or "previous_shot" (use previous shot's generated image)
2. If "previous_shot": are there entities (characters/props) that were visible in the previous shot but NOT in this shot?
   - If yes, list them in "entities_to_remove" with their descriptions
   - These will be erased from the previous shot's image before using it as background
3. silhouette_description: describe the dotted-line silhouettes to add (based on character count, approximate positions from T2I)
4. character_prompt: the T2I split for characters only (no background descriptions)
   - NEVER include skin tone/face color modifiers
   - Soften gore descriptions

Rules:
- Use "previous_shot" when the background state CHANGES due to story events (blood appears, room gets messy, items move)
- Use "base_set" when the shot shows the room in its default/original state
- First shot always uses "base_set"
- "entities_to_remove" only applies when background_source is "previous_shot"
- Output valid JSON"""

    shots_text = ""
    for i, s in enumerate(shots):
        chars_str = ", ".join(f"{c}({names.get(c[:3], '?')})" for c in s["chars"])
        # T2I 요약 (300자 제한 — 토큰 절약)
        t2i_short = s['t2i_prompt'][:300]
        shots_text += (
            f"### [{i}] {s['label']} | Characters: {chars_str} | Props: {s['props']}\n"
            f"Beat: {s['beat_title']}\n"
            f"T2I: {t2i_short}\n\n"
        )

    user = f"""## Shot sequence (same location, in order):
{shots_text}

## Output:
```json
{{
  "shots": [
    {{
      "shot_label": "S05-Shot1",
      "background_source": "base_set",
      "base_set_side": "IMAGE_1 or IMAGE_2",
      "entities_to_remove": [],
      "silhouette_description": "2 silhouettes: one near the sink (tall, standing), one by the doorway (medium, entering)",
      "character_prompt": "Characters/actions only, no background"
    }}
  ]
}}
```"""

    print("[Phase 1] LLM 전체 샷 시퀀스 분석...", flush=True)
    result = call_text(step="scene_director", system_prompt=system, user_prompt=user, temperature=0.1)
    text = result if isinstance(result, str) else ""
    # Debug: save raw response
    debug_dir = pathlib.Path(__file__).resolve().parent.parent.parent / "frontend" / "public" / "experiment" / "set_v8"
    debug_dir.mkdir(parents=True, exist_ok=True)
    with open(debug_dir / "llm_raw_response.txt", "w") as f:
        f.write(text)
    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"LLM returned empty JSON. Raw length: {len(result)}")
    return json.loads(text)


# ──────────────────────────────────────────────
# 이미지 생성 함수들
# ──────────────────────────────────────────────

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 generate_cleaned_bg(client, prev_shot_bytes, entities_to_remove):
    """이전 샷 이미지에서 지정된 요소를 제거한 배경 생성."""
    # entities_to_remove can be list of strings or list of dicts
    remove_descriptions = "\n".join(
        f"- {e}" if isinstance(e, str) else f"- {e.get('description', str(e))}"
        for e in entities_to_remove
    )

    prompt = (
        "Take the reference image and REMOVE the following elements from it. "
        "Replace them with the natural background that would be behind them. "
        "Keep everything else exactly the same.\n\n"
        f"REMOVE these items:\n{remove_descriptions}\n\n"
        "Output: the same scene but with those items erased, showing only the empty room/background. "
        "No people, no figures. Photorealistic."
    )

    labeled_refs = [("Source image — remove specified items from this:", prev_shot_bytes)]

    img_bytes, _ = client.generate_image(
        prompt=prompt, aspect_ratio="16:9", labeled_references=labeled_refs
    )
    return img_bytes


def generate_bg_with_silhouettes(client, bg_bytes, silhouette_desc):
    """배경에 점선 실루엣 추가."""
    prompt = (
        "Take the reference image (a room background) and add DOTTED LINE SILHOUETTES "
        "of people as placement guides. The silhouettes should be:\n"
        f"{silhouette_desc}\n\n"
        "Rules:\n"
        "- Silhouettes are dashed/dotted outlines only — no fill, no faces, no details\n"
        "- Semi-transparent, clearly visible but not dominant\n"
        "- Keep the background EXACTLY as in the reference image\n"
        "- No real people, only dotted outlines"
    )
    labeled_refs = [("Background to add silhouettes to:", bg_bytes)]

    img_bytes, _ = client.generate_image(
        prompt=prompt, aspect_ratio="16:9", labeled_references=labeled_refs
    )
    return img_bytes


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 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',
                  'shocked expression', text, flags=re.IGNORECASE)
    return re.sub(r'\s+', ' ', text).strip()


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

def main():
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    (OUT_DIR / "bg").mkdir(exist_ok=True)
    (OUT_DIR / "bg_silhouette").mkdir(exist_ok=True)
    (OUT_DIR / "shots").mkdir(exist_ok=True)

    shots, names, etm = load_data()
    print(f"{'='*60}")
    print(f"Set Design v8 (Chain + Silhouette) — L05 | {len(shots)} shots")
    print(f"{'='*60}\n")

    # Load v7 base images
    v7_dir = EXP / "set_v7"
    if not (v7_dir / "image_1.png").exists():
        print("ERROR: v7 base images not found. Run experiment_set_v7 first.")
        return
    base_images = {
        "IMAGE_1": (v7_dir / "image_1.png").read_bytes(),
        "IMAGE_2": (v7_dir / "image_2.png").read_bytes(),
    }
    print("Base images loaded from v7\n")

    # Phase 1: LLM analysis
    analysis = analyze_shot_sequence(shots, names)
    shot_plans = {s["shot_label"]: s for s in analysis["shots"]}

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

    print(f"\n샷별 계획:")
    for sp in analysis["shots"]:
        src = sp["background_source"]
        remove = len(sp.get("entities_to_remove", []))
        print(f"  {sp['shot_label']} → {src}" + (f" (remove {remove})" if remove else ""))
    print()

    # Phase 2: 순차 생성
    client = GeminiImageClient(model=settings.gemini_image_model)
    client.set_context(step="set_v8", operation_type="experiment",
                       project_id=PROJECT_ID, episode_id=EPISODE_ID)
    db = SessionLocal()

    prev_shot_image = None  # 이전 샷의 생성 이미지

    try:
        for i, shot in enumerate(shots):
            label = shot["label"]
            plan = shot_plans.get(label)
            if not plan:
                print(f"  [{label}] SKIP — no plan")
                continue

            src = plan["background_source"]
            char_prompt = sanitize_gore(plan["character_prompt"])
            sil_desc = plan.get("silhouette_description", "")
            entities_to_remove = plan.get("entities_to_remove", [])

            print(f"\n[{label}] source={src}", end="")

            # ── Step A: 배경 결정 ──
            if src == "base_set":
                side = plan.get("base_set_side", "IMAGE_1")
                bg_bytes = base_images.get(side, base_images["IMAGE_1"])
                print(f" ({side})")
            elif src == "previous_shot" and prev_shot_image:
                if entities_to_remove:
                    # 요소 제거 필요
                    print(f" → cleaning {len(entities_to_remove)} entities... ", end="", flush=True)
                    t0 = time.time()
                    try:
                        bg_bytes = generate_cleaned_bg(client, prev_shot_image, entities_to_remove)
                        (OUT_DIR / "bg" / f"{label}_cleaned.png").write_bytes(bg_bytes)
                        print(f"OK {time.time()-t0:.1f}s")
                    except Exception as e:
                        print(f"FAIL: {e} → fallback to base")
                        side = plan.get("base_set_side", "IMAGE_1")
                        bg_bytes = base_images.get(side, base_images["IMAGE_1"])
                else:
                    bg_bytes = prev_shot_image
                    print(" (prev shot direct)")
            else:
                side = plan.get("base_set_side", "IMAGE_1")
                bg_bytes = base_images.get(side, base_images["IMAGE_1"])
                print(f" (fallback {side})")

            # ── Step B: 점선 실루엣 추가 ──
            if sil_desc:
                print(f"  실루엣 추가... ", end="", flush=True)
                t0 = time.time()
                try:
                    bg_with_sil = generate_bg_with_silhouettes(client, bg_bytes, sil_desc)
                    (OUT_DIR / "bg_silhouette" / f"{label}_sil.png").write_bytes(bg_with_sil)
                    print(f"OK {time.time()-t0:.1f}s")
                except Exception as e:
                    print(f"FAIL: {e} → skip silhouette")
                    bg_with_sil = bg_bytes
            else:
                bg_with_sil = bg_bytes

            # ── Step C: 최종 샷 생성 ──
            entity_refs = get_composite_refs(db, shot["visible_entities"], char_prompt)
            labeled_refs = [("SET BACKGROUND with actor placement guides", bg_with_sil)]
            labeled_refs.extend(entity_refs)

            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"  샷 생성 (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
                )
                out_path = OUT_DIR / "shots" / f"{label.replace('-','_')}.png"
                out_path.write_bytes(img_bytes)
                prev_shot_image = img_bytes  # 다음 샷의 참조용
                print(f"OK {len(img_bytes)//1024}KB {time.time()-t0:.1f}s")
            except Exception as e:
                print(f"FAIL: {e}")
                # SAFETY fallback: base set 사용
                if "SAFETY" in str(e) or "moderation" in str(e):
                    print(f"  → parent fallback... ", end="", flush=True)
                    side = plan.get("base_set_side", "IMAGE_1")
                    labeled_refs[0] = ("SET BACKGROUND", base_images.get(side, base_images["IMAGE_1"]))
                    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
                        )
                        out_path = OUT_DIR / "shots" / f"{label.replace('-','_')}.png"
                        out_path.write_bytes(img_bytes)
                        prev_shot_image = img_bytes
                        print(f"OK (fallback)")
                    except Exception as e2:
                        print(f"FAIL: {e2}")

    finally:
        db.close()

    # URLs
    base_url = "http://192.168.231.91:3002/experiment/set_v8"
    print(f"\n{'='*60}")
    print("URLs")
    print(f"{'='*60}")
    print("\nCleaned BGs:")
    for p in sorted((OUT_DIR / "bg").glob("*.png")):
        print(f"  {base_url}/bg/{p.name}")
    print("\nSilhouette BGs:")
    for p in sorted((OUT_DIR / "bg_silhouette").glob("*.png")):
        print(f"  {base_url}/bg_silhouette/{p.name}")
    print("\nShots:")
    for p in sorted((OUT_DIR / "shots").glob("*.png")):
        print(f"  {base_url}/shots/{p.name}")


if __name__ == "__main__":
    main()