"""실험: 옥탑방 씬 도면(floor plan) 기반 공간/인물/카메라 가시화 — v2

변경점 vs v1 (초기):
- **GPT-5.5**가 씬/샷 원문 + staging + fixed_elements 전체를 보고 T2I 프롬프트 자동 생성
- 컨텍스트 절대 자르지 않음 (시나리오 전문 + 샷 description 원문 + camera_direction 전문 + fixed_elements 영문 description 원문)
- 옥탑방 관련 **11개 씬 + 7개 location** 모두 전달
- 생성된 프롬프트로 **gpt-image-2** 멀티턴 (base → edit 체인)

파이프라인:
  [씬/샷/staging/consistency/location 수집]
    → GPT-5.5 (JSON 구조)
    → { base_prompt, shot_prompts[] }
    → gpt-image-2 generate(base) + edit(shot별)

출력:
- backend/scripts/output/floor_plan/prompts.json — LLM 출력 + 원본 컨텍스트
- backend/scripts/output/floor_plan/base.png
- backend/scripts/output/floor_plan/S##_Shot#_*.png
"""
from __future__ import annotations

import base64
import json
import logging
import os
import sys
from pathlib import Path
from typing import Dict, Any, List

BACKEND = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(BACKEND))

from dotenv import load_dotenv  # noqa: E402
load_dotenv(BACKEND / ".env")

from openai import OpenAI  # noqa: E402

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("floor_plan")

PROJECT_ID = "d8a3b254-45c9-4eb4-bba6-8bae95eca96f"
EPISODE_ID = "e71c5a18-88e7-40e2-81bd-dc28c56a62f7"
OUT_DIR = BACKEND / "scripts" / "output" / "floor_plan"
OUT_DIR.mkdir(parents=True, exist_ok=True)

# 옥탑방 관련 범위
OKTAP_SCENES = [4, 5, 10, 11, 12, 13, 14, 17, 18, 25, 27]
OKTAP_LOCATIONS = ["L04", "L05", "L06", "L12", "L13", "L14", "L17"]

# 모델
TEXT_MODEL = "gpt-5.5"  # 조직 보유 최신 chat model
IMAGE_MODEL = "gpt-image-2"  # 2026-04-21 출시, thinking + 2K
IMAGE_SIZE = "1024x1024"
IMAGE_QUALITY = "high"

client = OpenAI()


# ──────────────────────────────────────────────────────────
# 데이터 로드 — 잘라내기 절대 금지
# ──────────────────────────────────────────────────────────

def load_cp(step_id: str) -> dict:
    base = BACKEND.parent / "projects" / PROJECT_ID / "checkpoints" / "episodes" / EPISODE_ID / step_id
    cur = base / "manifest.json"
    if cur.exists():
        return json.loads(cur.read_text(encoding="utf-8"))
    arches = sorted(base.glob("manifest_*.json"), reverse=True)
    if arches:
        logger.info("%s: using archive %s", step_id, arches[0].name)
        return json.loads(arches[0].read_text(encoding="utf-8"))
    raise FileNotFoundError(step_id)


def collect_full_context() -> Dict[str, Any]:
    """옥탑방 관련 모든 원문을 자르지 않고 수집."""
    ctx: Dict[str, Any] = {"scope": {"scenes": OKTAP_SCENES, "locations": OKTAP_LOCATIONS}}

    # 1) 씬 원문 — 전체
    scene_save = load_cp("scene_save")
    ctx["scene_texts"] = []
    for seg in scene_save["data"]["segments"]:
        if seg["scene_index"] in OKTAP_SCENES:
            ctx["scene_texts"].append({
                "scene_index": seg["scene_index"],
                "heading": seg["heading"],
                "text": seg["text"],  # 원문 전체, 자르기 금지
            })

    # 2) 샷 원문 — 옥탑방 씬의 모든 샷
    shot_cp = load_cp("shot_validator")
    ctx["shots"] = []
    for sc in shot_cp.get("data", {}).get("scenes", []):
        if sc.get("scene_index") in OKTAP_SCENES:
            for sh in sc.get("shots", []):
                ctx["shots"].append({
                    "scene_index": sc["scene_index"],
                    "shot_index": sh.get("shot_index"),
                    "description": sh.get("description", ""),  # 전문
                    "characters": sh.get("characters", []),
                    "based_on_beat": sh.get("based_on_beat"),
                })

    # 3) shot_staging — camera_direction 전문
    staging_cp = load_cp("shot_staging")
    ctx["staging"] = []
    for sh in staging_cp["data"]["shots"]:
        if sh.get("scene_index") in OKTAP_SCENES:
            ctx["staging"].append({
                "scene_index": sh["scene_index"],
                "shot_index": sh["shot_index"],
                "camera_direction": sh.get("camera_direction", ""),  # 전문
                "lighting_mood": sh.get("lighting_mood", ""),
                "perspective": sh.get("perspective", ""),
                "pov_character": sh.get("pov_character", ""),
                "key_bg_elements": sh.get("key_bg_elements", []),
                "character_angles": sh.get("character_angles", []),
            })

    # 4) scene_consistency fixed_elements — 전체
    sc_cp = load_cp("scene_consistency")
    ctx["fixed_elements"] = []
    for s in sc_cp["data"]["scenes"]:
        if s["scene_index"] in OKTAP_SCENES:
            ctx["fixed_elements"].append({
                "scene_index": s["scene_index"],
                "analysis_summary": s.get("analysis_summary", ""),
                "fixed_elements": s.get("fixed_elements", []),
            })

    # 5) scene_director — primary_location + present_entity_ids
    sd_cp = load_cp("scene_director")
    ctx["director"] = []
    for s in sd_cp["data"]["scenes"]:
        if s.get("scene_index") in OKTAP_SCENES:
            ctx["director"].append({
                "scene_index": s["scene_index"],
                "primary_location": s.get("primary_location"),
                "present_entity_ids": s.get("present_entity_ids", []),
            })

    # 6) locations 상세 — 7개 전체
    loc_cp = load_cp("entity_extract_location")
    locs = {l["short_id"]: l for l in loc_cp["data"]["locations"]}
    ctx["locations"] = []
    for lid in OKTAP_LOCATIONS:
        l = locs.get(lid)
        if l:
            ctx["locations"].append({
                "short_id": lid,
                "name": l.get("name"),
                "description": l.get("description", ""),
                "visual_traits": l.get("visual_traits", []),
            })

    # 7) entity_merge — 주요 인물/소품 맥락 (인물만 추려서)
    em_cp = load_cp("entity_merge")
    ctx["characters"] = [
        {"short_id": c.get("short_id"), "name": c.get("name"), "description": c.get("description", "")}
        for c in em_cp["data"].get("characters", [])
    ]

    return ctx


# ──────────────────────────────────────────────────────────
# GPT-5.5 — 도면 T2I 프롬프트 생성
# ──────────────────────────────────────────────────────────

SYSTEM_PROMPT = """당신은 범죄 수사 도면 작성 전문가입니다. 주어진 시나리오 씬/샷 데이터를 바탕으로 gpt-image-2로 생성할 **건축 도면(architectural floor plan) 스타일 T2I 프롬프트**를 영문으로 작성하세요.

## 출력 규칙
1. 프롬프트는 모두 **영문**. 단 공간/인물 라벨은 한국어 병기 가능 (예: "Living · 거실", "Suriyoung · 수리영").
2. **단일 건물 통합 도면** — 옥탑방 실내(거실·주방·민숙 방·수리영 방·안방) + 옥상 외부(옥탑 옥상) + 외곽(계단·마당) 모두 하나의 top-down 도면에 포함. 건물 경계와 외부 공간이 연속되게.
3. 스타일: 흑백 건축 도면 + 기술 도면 라인 드로잉. 벽=검정 실선, 문=swing arc, 창=이중 평행선, 가구=단순 외곽선 + 라벨, compass + scale bar.
4. **인물**: 살아있는 인물은 검정 원 + 방향 화살표, **사망 인물은 X-cross body outline + 점선 범죄현장 윤곽 + grey stippled 혈흔 pool**.
5. **카메라**: 검정 삼각형 wedge + dashed FOV cone. 라벨 형식 "CAM-S##-Shot# (WS/MS/CU, angle, 주석)".
6. **고정 요소 (fixed_elements)**: 씬 내 여러 샷에 걸친 고정 visual element는 **반드시 도면에 명시적으로 표기** (벽의 붉은 원, 쓰러진 의자, 커튼, 혈흔 등). scene_consistency의 영문 description을 그대로 반영.

## 출력 JSON 구조 (엄격)
```json
{
  "base_prompt": "string (인물 없음, 모든 공간 + 가구 + 소품 기본 상태)",
  "shot_prompts": [
    {
      "scene_index": 12,
      "shot_index": 4,
      "label": "S12_Shot4_murder_wide",
      "description": "한 줄 요약 (한국어)",
      "prompt": "string (base 위에 overlay할 annotation 지시. 인물 + 시체 + 혈흔 + 카메라 위치 + FOV cone + fixed_elements 반영)"
    }
  ]
}
```

## 어느 샷을 포함할지
사용자가 '시체 + 카메라 포함 도면'을 원했으므로 다음 우선순위:
- S12 (엄마 시신 발견) 모든 샷 — 시체 위치, 혈흔, 카메라
- S14 (시신 사라진 상태) 주요 샷 — 빈 공간 대조
- S25 (난장판 + 사진) 주요 샷
- 기타 옥탑방 씬은 선택적 (5개 이상 많지 않게)

최대 10개 샷 선정. 반드시 JSON만 출력 (추가 설명 금지).
"""


def gpt55_generate_prompts(ctx: Dict[str, Any]) -> Dict[str, Any]:
    logger.info("GPT-5.5 호출 — 컨텍스트 크기: scenes=%d shots=%d staging=%d fixed=%d",
                len(ctx["scene_texts"]), len(ctx["shots"]), len(ctx["staging"]), len(ctx["fixed_elements"]))

    user_content = (
        "## 옥탑방 관련 전체 컨텍스트 (금월도 1부)\n\n"
        f"### 씬 원문 ({len(ctx['scene_texts'])}개, 자르지 않음)\n"
        + "\n\n".join(f"【{s['heading']}】\n{s['text']}" for s in ctx["scene_texts"])
        + "\n\n"
        f"### 샷 description ({len(ctx['shots'])}개, 원문)\n"
        + "\n".join(
            f"- S{sh['scene_index']}_Shot{sh['shot_index']} (chars={sh['characters']}): {sh['description']}"
            for sh in ctx["shots"]
        )
        + "\n\n"
        f"### shot_staging ({len(ctx['staging'])}개)\n"
        + "\n".join(
            f"- S{st['scene_index']}_Shot{st['shot_index']}: camera={st['camera_direction']}; "
            f"lighting={st['lighting_mood']}; perspective={st['perspective']}; "
            f"key_bg={json.dumps(st['key_bg_elements'], ensure_ascii=False)}; "
            f"angles={json.dumps(st['character_angles'], ensure_ascii=False)}"
            for st in ctx["staging"]
        )
        + "\n\n"
        f"### scene_consistency fixed_elements ({len(ctx['fixed_elements'])}개)\n"
        + "\n".join(
            f"S{fe['scene_index']}: {fe['analysis_summary']}\n"
            + "\n".join(
                f"  - [{elem['element_type']}] {elem.get('element_id')} ({elem.get('character_name','')}): "
                f"{elem.get('description','')} — applies_to: {elem.get('applies_to_shots')}"
                for elem in fe["fixed_elements"]
            )
            for fe in ctx["fixed_elements"]
        )
        + "\n\n"
        f"### scene_director ({len(ctx['director'])}개)\n"
        + "\n".join(
            f"- S{d['scene_index']}: primary={d['primary_location']}, present={d['present_entity_ids']}"
            for d in ctx["director"]
        )
        + "\n\n"
        f"### 옥탑방 관련 locations ({len(ctx['locations'])}개)\n"
        + "\n".join(
            f"- {l['short_id']} ({l['name']}): {l['description']} | traits={l.get('visual_traits')}"
            for l in ctx["locations"]
        )
        + "\n\n"
        f"### 인물 목록 (맥락)\n"
        + "\n".join(
            f"- {c['short_id']} {c['name']}: {c['description']}"
            for c in ctx["characters"][:10]  # 주요 10명만 맥락용
        )
        + "\n\n"
        "이제 위 데이터를 바탕으로 JSON 출력하세요."
    )

    # chat.completions + json_object 모드
    resp = client.chat.completions.create(
        model=TEXT_MODEL,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_content},
        ],
        response_format={"type": "json_object"},
    )
    content = resp.choices[0].message.content
    data = json.loads(content)
    logger.info("GPT-5.5 결과: base_prompt 길이 %d, shot_prompts %d개",
                len(data.get("base_prompt", "")), len(data.get("shot_prompts", [])))
    return data


# ──────────────────────────────────────────────────────────
# gpt-image-2 — 멀티턴 도면 생성
# ──────────────────────────────────────────────────────────

def gen_image(prompt: str, out_path: Path) -> Path:
    logger.info("Generating base (%s, %s)...", IMAGE_MODEL, IMAGE_SIZE)
    resp = client.images.generate(
        model=IMAGE_MODEL, prompt=prompt, size=IMAGE_SIZE, quality=IMAGE_QUALITY, n=1,
    )
    out_path.write_bytes(base64.b64decode(resp.data[0].b64_json))
    logger.info("Saved: %s (%d KB)", out_path.name, out_path.stat().st_size // 1024)
    return out_path


def edit_image(base_path: Path, prompt: str, out_path: Path) -> Path:
    logger.info("Editing %s → %s", base_path.name, out_path.name)
    with open(base_path, "rb") as f:
        resp = client.images.edit(
            model=IMAGE_MODEL, image=f, prompt=prompt, size=IMAGE_SIZE, quality=IMAGE_QUALITY, n=1,
        )
    out_path.write_bytes(base64.b64decode(resp.data[0].b64_json))
    logger.info("Saved: %s (%d KB)", out_path.name, out_path.stat().st_size // 1024)
    return out_path


# ──────────────────────────────────────────────────────────
# 메인
# ──────────────────────────────────────────────────────────

def main() -> int:
    if not os.getenv("OPENAI_API_KEY"):
        logger.error("OPENAI_API_KEY not set")
        return 1

    # 1) 컨텍스트 수집
    ctx = collect_full_context()

    # 컨텍스트 원본 저장 (감사/재현용)
    (OUT_DIR / "context.json").write_text(
        json.dumps(ctx, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    logger.info("Context saved: %s (scenes=%d, shots=%d, staging=%d, fixed=%d, locations=%d)",
                OUT_DIR / "context.json",
                len(ctx["scene_texts"]), len(ctx["shots"]),
                len(ctx["staging"]), len(ctx["fixed_elements"]), len(ctx["locations"]))

    # 2) GPT-5.5로 프롬프트 생성
    prompts = gpt55_generate_prompts(ctx)
    (OUT_DIR / "prompts.json").write_text(
        json.dumps(prompts, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    logger.info("Prompts saved: prompts.json")

    # 3) base 도면 생성
    base_path = OUT_DIR / "base.png"
    gen_image(prompts["base_prompt"], base_path)

    # 4) 각 shot_prompt 멀티턴 edit
    for i, sp in enumerate(prompts.get("shot_prompts", []), 1):
        label = sp.get("label") or f"S{sp.get('scene_index')}_Shot{sp.get('shot_index')}"
        # 파일명 safe
        safe = "".join(c if c.isalnum() or c in ("_", "-") else "_" for c in label)
        out = OUT_DIR / f"{safe}.png"
        try:
            edit_image(base_path, sp["prompt"], out)
        except Exception as exc:
            logger.error("Failed %s: %s", label, exc)

    logger.info("=== DONE. Output: %s ===", OUT_DIR)
    return 0


if __name__ == "__main__":
    sys.exit(main())