"""실험: 세트 배경 재생성 (실루엣 제거) + 샷 재생성.

step1_compositions.json의 구도를 재사용하되 실루엣 지시를 제거.
체이닝으로 배경 4장 생성 → 샷 11장 생성.

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

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.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_ROOT = pathlib.Path(__file__).resolve().parent.parent.parent / "frontend" / "public" / "experiment"
SET_SRC = EXP_ROOT / "set_chain"  # step1/step3 JSON 소스
SET_OUT = EXP_ROOT / "set_v2"     # 새 배경
SHOT_OUT = EXP_ROOT / "set_shots_v2"  # 새 샷


def strip_silhouette(prompt: str) -> str:
    """실루엣 관련 문구 제거 + 'no people, empty room' 추가."""
    # Remove silhouette instruction
    prompt = re.sub(r'NO people\s*[—–-]\s*instead include.*?placement guides\.?', '', prompt, flags=re.DOTALL)
    prompt = prompt.strip().rstrip('.')
    prompt += ". Absolutely no people, no silhouettes, no figures. Empty room only."
    return prompt


def generate_set_images(compositions: list[dict]) -> dict[str, pathlib.Path]:
    """체이닝 방식으로 세트 배경 생성 (실루엣 없음)."""
    SET_OUT.mkdir(parents=True, exist_ok=True)
    client = GeminiImageClient(model=settings.gemini_image_model)
    client.set_context(step="set_v2_generate", operation_type="experiment")

    results = {}
    prev_images: list[bytes] = []

    for comp in compositions:
        comp_id = comp["id"]
        prompt = strip_silhouette(comp["t2i_prompt"])

        if prev_images:
            prompt += (
                "\n\nIMPORTANT: This is the SAME physical room/location as the reference image(s) below. "
                "Maintain the same wall colors, flooring, furniture style, window shape, and overall condition. "
                "Only the camera angle and framing should differ."
            )

        ref_label = f"(refs: {len(prev_images)})" if prev_images else "(no refs)"
        out_path = SET_OUT / f"{comp_id}.png"
        print(f"[SET] {comp_id} ({comp['name_ko']}) {ref_label}... ", end="", flush=True)

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

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

    return results


def get_ref_images_for_shot(db, shot_data: dict) -> list[tuple[str, bytes]]:
    """DB에서 샷의 인물/소품 참조 이미지 로드."""
    ve = shot_data.get("visible_entities", [])
    labeled_refs = []
    t2i_list = shot_data.get("t2i_variations", [])
    t2i_text = t2i_list[0].get("t2i_prompt", "") if t2i_list else ""

    for short_id in ve:
        if short_id.startswith("L"):
            continue

        entity = (
            db.query(Entity)
            .filter(Entity.project_id == PROJECT_ID, Entity.short_id == short_id)
            .first()
        )
        if not entity:
            continue

        eid = entity.id
        etype = entity.entity_type

        if etype == "character":
            pattern = rf'{short_id}(O\d{{2,3}})'
            m = re.search(pattern, t2i_text)
            if m:
                outfit_sid = m.group(1)
                outlook_entity = (
                    db.query(Entity)
                    .filter(Entity.project_id == PROJECT_ID, Entity.short_id == outfit_sid)
                    .first()
                )
                if outlook_entity:
                    composite = (
                        db.query(ImageAsset)
                        .filter(
                            ImageAsset.project_id == PROJECT_ID,
                            ImageAsset.asset_type == "reference",
                            ImageAsset.prompt_used.like(f"%composite:{eid}:{outlook_entity.id}%"),
                        )
                        .order_by(ImageAsset.is_primary.desc(), ImageAsset.created_at.desc())
                        .first()
                    )
                    if composite and pathlib.Path(composite.file_path).exists():
                        composite_id = f"{short_id}{outfit_sid}"
                        labeled_refs.append((f"{composite_id} — character wearing outfit", pathlib.Path(composite.file_path).read_bytes()))
                        continue

                    outlook_ref = (
                        db.query(ImageAsset)
                        .filter(
                            ImageAsset.project_id == PROJECT_ID,
                            ImageAsset.entity_id == outlook_entity.id,
                            ImageAsset.asset_type == "reference",
                            ImageAsset.is_primary == 1,
                        )
                        .first()
                    )
                    if outlook_ref and pathlib.Path(outlook_ref.file_path).exists():
                        labeled_refs.append(("outfit appearance", pathlib.Path(outlook_ref.file_path).read_bytes()))

            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"{short_id} — 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"{short_id} — prop/object", pathlib.Path(prop_ref.file_path).read_bytes()))

    return labeled_refs


def load_entity_text_map() -> dict[str, str]:
    etm = {}
    with open(BASE / "entity_merge" / "manifest.json") as f:
        data = json.load(f).get("data", {})
    for category in ["characters", "locations", "props"]:
        for e in data.get(category, []):
            sid = e.get("short_id", "")
            desc = e.get("short_description", e.get("name", ""))
            if sid and desc:
                etm[sid] = desc
    with open(BASE / "entity_t2i" / "manifest.json") as f:
        t2i_data = json.load(f).get("data", {}).get("completed", {})
    for name, info in t2i_data.items():
        sid = info.get("short_id", "")
        desc = info.get("short_description", info.get("description", ""))
        if sid and desc and sid not in etm:
            etm[sid] = desc
    return etm


def main():
    # Load existing compositions & assignments
    with open(SET_SRC / "step1_compositions.json") as f:
        compositions = json.load(f)["compositions"]
    with open(SET_SRC / "step3_assignments.json") as f:
        assignments = json.load(f)["assignments"]

    with open(BASE / "scene_detail" / "manifest.json") as f:
        detail_data = json.load(f)
    shot_details = {}
    for s in detail_data["data"]["scenes"]:
        if "L05" not in s.get("visible_entities", []):
            continue
        label = f"S{s['scene_index']:02d}-Shot{s.get('_shot_index', 1)}"
        shot_details[label] = s

    entity_text_map = load_entity_text_map()

    # ── Step 1: 세트 배경 재생성 (실루엣 제거) ──
    print("=== 세트 배경 재생성 (실루엣 제거) ===\n")
    set_images_paths = generate_set_images(compositions)
    set_images = {cid: p.read_bytes() for cid, p in set_images_paths.items()}
    print()

    # ── Step 2: 샷 이미지 재생성 ──
    print("=== 샷 이미지 재생성 ===\n")
    SHOT_OUT.mkdir(parents=True, exist_ok=True)

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

    db = SessionLocal()
    try:
        for assign in assignments:
            shot_label = assign["shot_label"]
            comp_id = assign["composition_id"]
            shot_data = shot_details.get(shot_label)
            if not shot_data:
                continue

            set_img = set_images.get(comp_id)
            if not set_img:
                continue

            t2i_list = shot_data.get("t2i_variations", [])
            if not t2i_list:
                continue

            t2i_prompt = t2i_list[0]["t2i_prompt"]
            entity_refs = get_ref_images_for_shot(db, shot_data)

            labeled_refs = [
                ("SET BACKGROUND — use this exact room as the background environment", set_img),
            ]
            labeled_refs.extend(entity_refs)

            print(f"[{shot_label}] → {comp_id} | refs: 1+{len(entity_refs)} | ", end="", flush=True)

            try:
                final_prompt = _build_final_scene_prompt(
                    t2i_prompt=t2i_prompt,
                    labeled_refs=labeled_refs,
                    style_context="",
                    scene_index=shot_data.get("scene_index", 0),
                    entity_text_map=entity_text_map,
                )
            except Exception as e:
                print(f"PROMPT FAIL: {e}")
                continue

            t0 = time.time()
            try:
                img_bytes, _ = client.generate_image(
                    prompt=final_prompt,
                    aspect_ratio="16:9",
                    labeled_references=labeled_refs,
                )
                filename = shot_label.replace("-", "_") + ".png"
                (SHOT_OUT / 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()

    print("\n=== URLs ===")
    print("\n세트 배경:")
    for cid in ["SET_A", "SET_B", "SET_C", "SET_D"]:
        print(f"  http://192.168.231.91:3002/experiment/set_v2/{cid}.png")
    print("\n샷:")
    for p in sorted(SHOT_OUT.glob("S*_Shot*.png")):
        print(f"  http://192.168.231.91:3002/experiment/set_shots_v2/{p.name}")


if __name__ == "__main__":
    main()