"""실험: 세트 배경 참조 이미지를 넣어서 샷별 이미지 생성.

- set_chain/ 의 SET_A~D 이미지를 배경 참조로 사용
- step3_assignments.json의 배정에 따라 각 샷에 맞는 세트 이미지 첨부
- 기존 _build_final_scene_prompt를 import해서 래핑 (원본 코드 수정 없음)
- 출력: frontend/public/experiment/set_shots/S05_Shot1.png 형태

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

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.modules.llm.gemini_image_client import GeminiImageClient
from app.services.prompt_service import build_final_scene_prompt as _build_final_scene_prompt
from app.core.database import SessionLocal
from app.models.project import ImageAsset, EntityCanon as Entity, Episode

# ── 경로 ──
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}"
)
SET_DIR = pathlib.Path(__file__).resolve().parent.parent.parent / "frontend" / "public" / "experiment" / "set_chain"
OUT_DIR = pathlib.Path(__file__).resolve().parent.parent.parent / "frontend" / "public" / "experiment" / "set_shots"


def load_assignments() -> list[dict]:
    with open(SET_DIR / "step3_assignments.json") as f:
        return json.load(f)["assignments"]


def load_set_images() -> dict[str, bytes]:
    """SET_A.png ~ SET_D.png 로드."""
    images = {}
    for p in SET_DIR.glob("SET_*.png"):
        images[p.stem] = p.read_bytes()
    return images


def load_scene_detail() -> dict[str, dict]:
    """scene_detail에서 L05 포함 샷의 상세 정보 로드. key = shot_label."""
    with open(BASE / "scene_detail" / "manifest.json") as f:
        data = json.load(f)

    shots = {}
    for s in data["data"]["scenes"]:
        ve = s.get("visible_entities", [])
        if "L05" not in ve:
            continue
        si = s["scene_index"]
        shi = s.get("_shot_index", 1)
        label = f"S{si:02d}-Shot{shi}"
        shots[label] = s
    return shots


def get_ref_images_for_shot(db, shot_data: dict) -> list[tuple[str, bytes]]:
    """DB에서 해당 샷의 visible_entities에 대한 참조 이미지(인물/소품만) 로드."""
    ve = shot_data.get("visible_entities", [])
    labeled_refs = []

    for short_id in ve:
        if short_id.startswith("L"):
            continue  # 배경은 세트 이미지로 대체

        # short_id → entity UUID 찾기
        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

        # composite 이미지 우선 (C##O## 패턴)
        # T2I에서 C##O## 패턴 추출
        t2i_list = shot_data.get("t2i_variations", [])
        t2i_text = t2i_list[0].get("t2i_prompt", "") if t2i_list else ""

        if etype == "character":
            # T2I에서 이 캐릭터의 C##O## 패턴 찾기
            pattern = rf'{short_id}(O\d{{2,3}})'
            m = re.search(pattern, t2i_text)
            if m:
                outfit_sid = m.group(1)
                # 아웃룩 entity 찾기
                outlook_entity = (
                    db.query(Entity)
                    .filter(Entity.project_id == PROJECT_ID, Entity.short_id == outfit_sid)
                    .first()
                )
                if outlook_entity:
                    # composite 이미지 찾기
                    composite = (
                        db.query(ImageAsset)
                        .filter(
                            ImageAsset.project_id == PROJECT_ID,
                            ImageAsset.asset_type == "reference",
                            ImageAsset.prompt_used.like(f"%composite:{eid}:{outlook_entity.id}%"),
                            ImageAsset.is_primary == 1,
                        )
                        .first()
                    )
                    if composite and pathlib.Path(composite.file_path).exists():
                        composite_id = f"{short_id}{outfit_sid}"
                        label = f"{composite_id} — character wearing outfit"
                        labeled_refs.append((label, pathlib.Path(composite.file_path).read_bytes()))
                        continue

                    # composite 없으면 아웃룩 단독
                    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]:
    """entity_detail에서 short_id → 텍스트 설명 맵 구성."""
    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

    # entity_t2i에서도 보완
    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():
    OUT_DIR.mkdir(parents=True, exist_ok=True)

    assignments = load_assignments()
    set_images = load_set_images()
    shot_details = load_scene_detail()
    entity_text_map = load_entity_text_map()

    print(f"=== Set Shot Generation ===")
    print(f"Assignments: {len(assignments)}")
    print(f"Set images: {list(set_images.keys())}")
    print(f"Entity text map: {len(entity_text_map)} entries")
    print()

    client = GeminiImageClient(model=settings.gemini_image_model)
    client.set_context(
        step="set_shot_experiment",
        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:
                print(f"[{shot_label}] SKIP — no shot data")
                continue

            set_img = set_images.get(comp_id)
            if not set_img:
                print(f"[{shot_label}] SKIP — no set image for {comp_id}")
                continue

            t2i_list = shot_data.get("t2i_variations", [])
            if not t2i_list:
                print(f"[{shot_label}] SKIP — no T2I prompt")
                continue

            t2i_prompt = t2i_list[0]["t2i_prompt"]

            # 참조 이미지 구성: 세트 배경 + 인물/소품
            entity_refs = get_ref_images_for_shot(db, shot_data)

            # 세트 배경을 첫 번째 참조로 추가
            labeled_refs = [
                (f"SET BACKGROUND — use this exact room as the background environment", set_img),
            ]
            labeled_refs.extend(entity_refs)

            # _build_final_scene_prompt로 최종 프롬프트 구성
            print(f"[{shot_label}] → {comp_id} | refs: 1(set) + {len(entity_refs)}(entity) | ", 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, resp_ms = client.generate_image(
                    prompt=final_prompt,
                    aspect_ratio="16:9",
                    labeled_references=labeled_refs,
                )
                # S05-Shot1 → S05_Shot1.png
                filename = shot_label.replace("-", "_") + ".png"
                out_path = OUT_DIR / filename
                out_path.write_bytes(img_bytes)
                elapsed = time.time() - t0
                print(f"OK  {len(img_bytes)//1024}KB  {elapsed:.1f}s")
            except Exception as e:
                elapsed = time.time() - t0
                print(f"FAIL ({elapsed:.1f}s): {e}")

    finally:
        db.close()

    # 결과 요약
    print()
    print("=== Results ===")
    generated = sorted(OUT_DIR.glob("S*_Shot*.png"))
    for p in generated:
        print(f"  http://192.168.231.91:3002/experiment/set_shots/{p.name}")


if __name__ == "__main__":
    main()