"""실험 v3 — Multi-step floor plan reasoning (시나리오 무관)

테스팅 전용. 메인 파이프라인 통합 X. 시스템 프롬프트는 도면 방법론만 (시나리오 의존 어휘 절대 금지).

## 흐름
Step 1 — 공간 분석 (LLM 호출 1)
  입력: 관련 씬 원문 + 샷 + fixed_elements + location description + director_notes
  출력: space_groups (도면 분할 계획) + shared_elements

Step 2 — Base 도면 (LLM 호출 1, 모든 plan 한번에)
  입력: Step 1 결과
  출력: 각 space_group별 T2I 프롬프트 + 범례 + 요소 메타

Step 3 — 샷별 Overlay (LLM 호출 N, 선택 샷마다)
  입력: Step 1+2 결과 + 해당 샷 데이터
  출력: 카메라 좌표 + 인물 자세/위치 + 추가 요소 + T2I 프롬프트

## 인자 (CLI)
--project-id, --episode-id, --scene-indices, --location-ids, --target-shots scene:shot ...

## 출력
output/floor_plan_test/<run_id>/
  context.json
  step1_spatial.json
  step2_base_plans.json
  step3_shot_<S##_Shot#>.json
  base_<plan_id>.png
  shot_<S##_Shot#>.png
  manifest.json
"""
from __future__ import annotations

import argparse
import base64
import json
import logging
import os
import sys
import time
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_v3")


# ──────────────────────────────────────────────────────────
# 시스템 프롬프트 — 도면 방법론만 (시나리오 어휘 0)
# ──────────────────────────────────────────────────────────

SYSTEM_STEP1 = """당신은 건축 도면 분석가입니다. 영화 시나리오 씬/샷 데이터를 받아 **공간 구조를 분석**하고 도면 분할 계획을 수립합니다.

## 분석 원칙
1. **데이터 의존만**: 사용자 메시지의 씬 원문 / 샷 description / location description / fixed_elements에 **명시된 요소만** 사용. 추측·창작 금지.
2. **공간 그룹핑**: 한 평면 도면에 모두 담기 어려우면 분할. 기준:
   - 실내/실외 차이 큼 → 분할
   - 층/높이 차이 (예: 옥상 vs 1층 마당) → 분할
   - 스케일 차이 10배 이상 (작은 방 vs 마당 전체) → 분할
   - 같은 층 같은 스케일은 통합
3. **shared_elements**: 여러 씬에 걸쳐 동일하게 유지되는 시각 요소 (fixed_elements 영문 description 활용).
4. **scale_estimate**: 데이터의 단서로 합리적 추정. 정확한 평수 단정 금지 — `"approx Nm x Mm"` 모호 표기.
5. **rooms / sub-areas**: 각 공간 내부 세부 구획. 벽 / 문 / 창문 / 가구 위치는 데이터에서 직접 추출.

## 출력 JSON (엄격)
```
{
  "space_groups": [
    {
      "id": "string (snake_case unique id)",
      "label": "string (Korean · English bilingual)",
      "covers_location_ids": ["L##", ...],
      "scale_estimate": "approx Nm x Mm",
      "type": "interior | exterior | site_map",
      "rooms": [
        {
          "name": "string (Korean optional · English)",
          "walls": ["description string"],
          "doors": [{"connects_to":"string","note":"string"}],
          "windows": [{"position":"string","note":"string"}],
          "furniture": [{"name":"string","position":"string","note":"string"}],
          "fixtures": [{"name":"string","position":"string"}]
        }
      ]
    }
  ],
  "shared_elements": [
    {
      "id": "string",
      "appears_in_scenes": [int],
      "in_space_group_id": "string",
      "in_room": "string",
      "description": "string (영문)",
      "note": "string (한국어 보조)"
    }
  ]
}
```

## ⚠ 라벨·id 규칙 — 고유명사 절대 금지 (CLAUDE.md 절대 규칙)
- `id` (snake_case): **일반 명사 기반만**. 인명·도시명·지역명·국가명·작품명 어떤 고유명사도 포함 금지.
  - ✅ "rooftop_room_interior", "residential_yard_and_stairs", "rooftop_site_map"
  - ❌ 어떤 도시·지역·인명도 id에 절대 금지 (입력 데이터에 그런 명사가 있어도 일반 명사로 변환)
- `label` 한영 병기 가능. 단 **인명·고유 지명·작품 명칭 제외**, 일반 명사만 (예: "옥탑방 내부 · rooftop room interior", "다세대 빌라 마당 · multi-family villa yard").
  - 입력 데이터에 "인천 변두리 다세대 빌라" 같은 표현이 있어도 → "다세대 빌라" 부분만 유지하고 지명("인천 변두리") 제거.
- `rooms[].name`, `walls`, `furniture`, `fixtures`, `note` 등 **모든 텍스트 필드에 인명·지명 등 고유명사 금지**. 일반 명사로 추상화.
- 시나리오에 명시되지 않은 가구/공간 추가 절대 금지.

JSON만 출력. 추가 설명 금지.
"""

SYSTEM_STEP2 = """당신은 건축 도면 작성 전문가입니다. Step 1의 spatial_analysis를 받아 각 space_group에 대해 **건축 도면 T2I 프롬프트**와 **요소 메타데이터**를 작성합니다.

## 도면 스타일 (모든 도면 공통, 시나리오 무관)
- Black-and-white architectural plan, top-down orthographic view, technical line drawing on white background
- Walls = solid black lines, doors = swing arcs, windows = double parallel lines
- Furniture = labeled simple geometric outlines
- Bilingual labels (Korean · English) **on the drawing only** (그래픽 라벨 한정)
- Compass rose (N) at corner, scale bar (e.g. 0-5m) at corner
- No 3D perspective, strictly 2D plan view
- No human figures in base plan (empty of characters)
- White background, no shading

## ⚠ T2I 프롬프트 본문 vs 도면 그래픽 라벨 분리
- **그래픽 라벨**(도면 안에 그려지는 텍스트, 일반 명사): 한국어 일반 명사 OK. t2i_prompt 안에서 `Label rooms with bilingual text "거실 · Living", "주방 · Kitchen"` 식으로 지시 가능. 단 **인명·작품 고유명사가 라벨에 들어가지 않도록** 일반 명사 사용 (예: "안방 · Master Bedroom", "작은방 · Small Bedroom").
- **t2i_prompt 산문**(공간 묘사, 가구 배치 설명): **인명·지명·작품 명칭 등 고유명사 절대 금지**. 보통명사만.
  - ❌ 어떤 캐릭터 이름·도시명·작품명도 t2i 본문에 등장 금지
  - ✅ "a small bedroom", "the bedroom wall", "a small rooftop residential apartment"

## 데이터 의존
- Step 1의 rooms / furniture / doors / windows / fixtures를 모두 도면에 표기
- 임의 추가 금지
- 입력 메타에 인명·지명이 새어나와 있더라도 **t2i_prompt 및 모든 출력 필드에서 일반 명사로 변환** (CLAUDE.md 절대 규칙)

## 출력 JSON (엄격)
```
{
  "base_plans": [
    {
      "id": "string (Step 1의 space_group.id 그대로)",
      "label": "string",
      "t2i_prompt": "string (영문, gpt-image-2가 그릴 도면 지시)",
      "legend": [{"symbol":"●","meaning":"..."}],
      "elements_meta": [
        {"id":"string","kind":"wall|door|window|furniture|fixture","label":"string","position_hint":"string","note":"string"}
      ],
      "annotations": [{"text":"string","placement":"string"}]
    }
  ]
}
```

JSON만 출력. 추가 설명 금지.
"""

SYSTEM_STEP3 = """당신은 영화 카메라 도면 분석가입니다. 한 샷의 카메라 위치 / 인물 자세 / 추가 요소를 base 도면 위에 overlay할 수 있도록 정밀 분석합니다.

## 카메라 표기 규칙
- triangular wedge symbol pointing in lens direction
- dashed FOV cone (wider for WS, narrower for CU)
- 라벨: `CAM-S##-Shot# (WS/MS/CU/MCU/ECU, angle, lens height, optional handheld/dolly note)`

## 인물·요소 표기 규칙
- 살아있는 인물: filled circle + 작은 방향 화살표(facing).
- 사망/움직이지 않는 인물: outline marker(X-cross 또는 chalk-style outline) + dashed contour
- 가구 변경(쓰러진 의자 등): 변경된 상태 그대로 그림
- 새 요소(흙발자국/벽 마크 등): grey/muddy 텍스처 영역 또는 explicit annotation

## ⚠ T2I 프롬프트 본문 — 고유명사 절대 금지 (CLAUDE.md 절대 규칙)
- t2i_prompt 산문에는 **인명·공간 고유명·지역명·작품 명칭 등 어떤 고유명사도 넣지 않는다**.
  - ❌ 인물 고유명, 캐릭터 이름 (어떤 언어든)
  - ❌ 도시·지역·국가 등 고유 지명
  - ❌ 작품 제목·세계관 명칭
  - ✅ 보통명사: "a small bedroom", "an adult human marker", "a low-rise residential rooftop"
- 인물 마커는 short_id(C##)로만 참조하거나 보통명사("adult human marker", "young person marker") 사용.
- 공간 묘사도 보통명사: "the kitchen-living room area", "the small bedroom", "the rooftop exterior"
- characters 필드의 `name` 한국어 보존은 **메타 데이터 기록용** (후속 파이프라인 참조). 절대 t2i_prompt에 직접 삽입 금지.

## 안전 어휘 (이미지 생성 moderation 회피, 도면 도메인이라도 단어 자체 회피)
- 사망 묘사: "motionless seated figure marker", "still figure outline" (avoid: dead, corpse, deceased, victim, body)
- 혈흔: "dark fluid stippled region", "reddish marker stain" (avoid: blood, gore, hemorrhage)
- 상처/외상: "condition note", "state indicator" (avoid: wound, trauma, torn flesh, injury)
- 범죄현장: "investigation marker zone" (avoid: crime scene, murder, homicide)
- 행위 동사 회피 (서사 동작도 도면 라벨에 직접 노출 금지):
  - avoid: kill, stab, strangle, choke, shoot, attack, assault, beat, slash, bite, rip
  - 사용: "altercation marker", "struggle indicator", "contact zone", "approach trail" 등 중립 용어
- 도구 명시 회피:
  - avoid: knife, blade, weapon, gun
  - 사용: "object marker (unspecified)" 또는 도면 범례에 형태(triangular/elongated)만 기재

## 출력 JSON (엄격)
```
{
  "scene_index": int,
  "shot_index": int,
  "base_plan_id": "string (Step 2 base_plans 중 하나)",
  "camera": {
    "position":"string (좌표 묘사)",
    "heading":"string (lens 방향)",
    "height":"low|eye-level|high|overhead",
    "fov":"WS|MS|MCU|CU|ECU",
    "lens_note":"string (handheld, dolly, slow push-in 등)"
  },
  "characters": [
    {"id":"C##","name":"string","status":"alive|motionless|other","posture":"string","position":"string","facing":"string"}
  ],
  "additions": [
    {"type":"string","position":"string","note":"string"}
  ],
  "t2i_prompt": "string (영문, base 도면 위에 overlay할 카메라/인물/요소 지시. base 가구는 유지하고 위 정보만 추가)",
  "legend_updates": [{"symbol":"","meaning":""}]
}
```

JSON만 출력. 추가 설명 금지.
"""


# ──────────────────────────────────────────────────────────
# 데이터 수집 (자르기 금지)
# ──────────────────────────────────────────────────────────

def cp_path(project_id: str, episode_id: str, step_id: str) -> Path:
    return BACKEND.parent / "projects" / project_id / "checkpoints" / "episodes" / episode_id / step_id


def load_cp(project_id: str, episode_id: str, step_id: str) -> dict:
    base = cp_path(project_id, 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:
        return json.loads(arches[0].read_text(encoding="utf-8"))
    raise FileNotFoundError(f"{step_id} (project={project_id} episode={episode_id})")


def collect_context(project_id: str, episode_id: str,
                    scene_indices: List[int], location_ids: List[str]) -> Dict[str, Any]:
    ctx: Dict[str, Any] = {
        "project_id": project_id,
        "episode_id": episode_id,
        "scope": {"scenes": scene_indices, "locations": location_ids},
    }

    # 씬 원문
    ss = load_cp(project_id, episode_id, "scene_save")
    ctx["scenes"] = [
        {"scene_index": s["scene_index"], "heading": s["heading"], "text": s["text"]}
        for s in ss["data"]["segments"] if s["scene_index"] in scene_indices
    ]

    # 샷 원문
    try:
        shv = load_cp(project_id, episode_id, "shot_validator")
    except FileNotFoundError:
        shv = load_cp(project_id, episode_id, "shot_extract")
    ctx["shots"] = []
    for sc in shv["data"].get("scenes", []):
        if sc.get("scene_index") in scene_indices:
            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", []),
                })

    # shot_staging
    st = load_cp(project_id, episode_id, "shot_staging")
    ctx["staging"] = [
        {
            "scene_index": s["scene_index"],
            "shot_index": s["shot_index"],
            "camera_direction": s.get("camera_direction", ""),
            "lighting_mood": s.get("lighting_mood", ""),
            "perspective": s.get("perspective", ""),
            "key_bg_elements": s.get("key_bg_elements", []),
            "character_angles": s.get("character_angles", []),
        }
        for s in st["data"].get("shots", []) if s.get("scene_index") in scene_indices
    ]

    # scene_consistency
    sc = load_cp(project_id, episode_id, "scene_consistency")
    ctx["fixed_elements"] = [
        {"scene_index": s["scene_index"],
         "analysis_summary": s.get("analysis_summary", ""),
         "fixed_elements": s.get("fixed_elements", [])}
        for s in sc["data"].get("scenes", []) if s["scene_index"] in scene_indices
    ]

    # locations — 누락 id 경고 (NIT-2)
    loc = load_cp(project_id, episode_id, "entity_extract_location")
    locs = {l["short_id"]: l for l in loc["data"].get("locations", [])}
    missing_loc_ids = [lid for lid in location_ids if lid not in locs]
    if missing_loc_ids:
        logger.warning("location_ids 누락(checkpoint에 없음): %s", missing_loc_ids)
    ctx["locations"] = [
        {"short_id": lid, "name": locs[lid].get("name"),
         "description": locs[lid].get("description", ""),
         "visual_traits": locs[lid].get("visual_traits", [])}
        for lid in location_ids if lid in locs
    ]

    # entities (characters) — 라벨용
    em = load_cp(project_id, episode_id, "entity_merge")
    ctx["characters"] = [
        {"short_id": c.get("short_id"), "name": c.get("name"),
         "description": c.get("description", "")}
        for c in em["data"].get("characters", [])
    ]

    # visual_world_rules — director_notes
    try:
        vwr = load_cp(project_id, episode_id, "visual_world_rules")
        ctx["visual_world"] = {
            "era": vwr["data"].get("era"),
            "region": vwr["data"].get("region"),
            "director_notes": vwr["data"].get("director_notes", []),
        }
    except FileNotFoundError:
        ctx["visual_world"] = {}

    return ctx


# ──────────────────────────────────────────────────────────
# 사용자 메시지 빌더
# ──────────────────────────────────────────────────────────

def msg_step1(ctx: Dict[str, Any]) -> str:
    return (
        "## 입력 데이터 (자르지 않음, 도면 분석용)\n\n"
        f"### 씬 원문 ({len(ctx['scenes'])}개)\n"
        + "\n\n".join(f"【{s['heading']}】\n{s['text']}" for s in ctx["scenes"])
        + "\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{s['scene_index']}_Shot{s['shot_index']}: cam={s['camera_direction']} | "
            f"perspective={s['perspective']} | key_bg={json.dumps(s['key_bg_elements'], ensure_ascii=False)} | "
            f"angles={json.dumps(s['character_angles'], ensure_ascii=False)}"
            for s in ctx["staging"]
        )
        + "\n\n"
        f"### fixed_elements ({len(ctx['fixed_elements'])}개)\n"
        + "\n".join(
            f"S{fe['scene_index']}: {fe['analysis_summary']}\n"
            + "\n".join(
                f"  - [{e['element_type']}] {e.get('element_id')} ({e.get('character_name','')}): "
                f"{e.get('description','')} | applies_to: {e.get('applies_to_shots')}"
                for e in fe["fixed_elements"]
            )
            for fe in ctx["fixed_elements"]
        )
        + "\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"### visual_world (참고)\n"
        f"era={ctx.get('visual_world', {}).get('era')}, region={ctx.get('visual_world', {}).get('region')}, "
        f"director_notes={ctx.get('visual_world', {}).get('director_notes', [])}\n\n"
        "위 데이터로 Step 1 spatial_analysis JSON을 출력하세요."
    )


def msg_step2(ctx: Dict[str, Any], spatial: Dict[str, Any]) -> str:
    """Step 2: spatial JSON + locations 원문 + 씬 헤딩 전달 (한국어 라벨 보존, HIGH-2 fix)."""
    return (
        "## Step 1 spatial_analysis 결과 (자르지 않음)\n```json\n"
        + json.dumps(spatial, ensure_ascii=False, indent=2)
        + "\n```\n\n"
        "## 원본 locations description (한국어 라벨 정확도용, 자르지 않음)\n"
        + "\n".join(
            f"- {l['short_id']} ({l['name']}): {l['description']} | traits={l.get('visual_traits')}"
            for l in ctx["locations"]
        )
        + "\n\n"
        "## 관련 씬 헤딩 (공간 맥락)\n"
        + "\n".join(f"- S{s['scene_index']}: {s['heading']}" for s in ctx["scenes"])
        + "\n\n"
        "위 데이터로 Step 2 base_plans JSON을 출력하세요. "
        "각 space_group마다 하나의 base plan. 도면 스타일 시스템 규칙 엄격 준수. "
        "T2I 프롬프트는 영문, 라벨은 한영 병기 (locations description의 한국어 명사 그대로 사용)."
    )


def msg_step3(ctx: Dict[str, Any], spatial: Dict[str, Any], plans: Dict[str, Any],
              scene_index: int, shot_index: int) -> str:
    # 해당 샷 데이터 추출
    shot_desc = next(
        (sh for sh in ctx["shots"]
         if sh["scene_index"] == scene_index and sh["shot_index"] == shot_index),
        None,
    )
    staging = next(
        (s for s in ctx["staging"]
         if s["scene_index"] == scene_index and s["shot_index"] == shot_index),
        None,
    )
    fixed = next((fe for fe in ctx["fixed_elements"] if fe["scene_index"] == scene_index), None)
    scene_text = next((s for s in ctx["scenes"] if s["scene_index"] == scene_index), None)

    # applies_to_shots 타입 정규화 (MEDIUM-1 fix)
    def _matches_shot(applies_to: Any) -> bool:
        for x in (applies_to or []):
            try:
                if int(x) == int(shot_index):
                    return True
            except (TypeError, ValueError):
                continue
        return False

    return (
        f"## 대상 샷: S{scene_index}_Shot{shot_index}\n\n"
        "### 해당 씬 원문\n"
        + (f"【{scene_text['heading']}】\n{scene_text['text']}\n\n" if scene_text else "")
        + "### 샷 description\n"
        + (f"{shot_desc['description']}\n  characters: {shot_desc['characters']}\n\n" if shot_desc else "(없음)\n\n")
        + "### shot_staging\n"
        + (json.dumps(staging, ensure_ascii=False, indent=2) + "\n\n" if staging else "(없음)\n\n")
        + "### fixed_elements (해당 씬 내, 이 샷에 applies_to 포함된 것)\n"
        + (
            "\n".join(
                f"  - [{e['element_type']}] {e.get('element_id')} ({e.get('character_name','')}): "
                f"{e.get('description','')} | applies_to: {e.get('applies_to_shots')}"
                for e in (fixed["fixed_elements"] if fixed else [])
                if _matches_shot(e.get("applies_to_shots"))
            ) or "(없음)"
        )
        + "\n\n"
        + "### Step 1 spatial_analysis (전체, 자르지 않음)\n```json\n"
        + json.dumps(spatial, ensure_ascii=False)  # 자르기 금지 (CLAUDE.md 절대 규칙)
        + "\n```\n\n"
        + "### Step 2 base_plans 요약 (id 매칭용)\n"
        + "\n".join(
            f"- {p['id']}: {p.get('label','')} (covers via spatial)"
            for p in plans.get("base_plans", [])
        )
        + "\n\n"
        f"위 데이터로 S{scene_index}_Shot{shot_index}의 overlay JSON을 출력하세요. "
        "**reasoning step**: (a) 어느 base_plan_id 위에 overlay할지 spatial.space_groups의 covers_location_ids와 해당 씬 primary_location 매칭으로 결정 → "
        "(b) camera 좌표를 도면 좌표계로 환산 → (c) 인물 status/posture/position을 데이터에서 직접 추출 (추측 금지) → "
        "(d) additions에 fixed_elements + 샷 고유 변경 반영 → (e) t2i_prompt 영문 작성. "
        "안전 어휘 규칙 엄격 준수 (사망/혈흔/외상 직접 단어 금지)."
    )


# ──────────────────────────────────────────────────────────
# LLM / Image
# ──────────────────────────────────────────────────────────

def call_llm_json(client: OpenAI, model: str, system: str, user: str) -> Dict[str, Any]:
    resp = client.chat.completions.create(
        model=model,
        messages=[{"role": "system", "content": system}, {"role": "user", "content": user}],
        response_format={"type": "json_object"},
    )
    return json.loads(resp.choices[0].message.content)


def _decode_or_raise(resp: Any, op: str) -> bytes:
    """b64_json None 방어 (MEDIUM-2 + NIT-3 fix)."""
    if not resp.data:
        raise RuntimeError(f"{op}: empty response.data")
    item = resp.data[0]
    b64 = getattr(item, "b64_json", None)
    if not b64:
        url = getattr(item, "url", None)
        raise RuntimeError(
            f"{op}: b64_json missing (url={url}). moderation 거부 또는 response_format mismatch 가능"
        )
    return base64.b64decode(b64)


def gen_image(client: OpenAI, model: str, prompt: str, size: str, quality: str, out_path: Path) -> Path:
    logger.info("[image] %s → %s", model, out_path.name)
    resp = client.images.generate(
        model=model, prompt=prompt, size=size, quality=quality, n=1,
    )
    out_path.write_bytes(_decode_or_raise(resp, f"gen_image({out_path.name})"))
    logger.info("  saved %d KB", out_path.stat().st_size // 1024)
    return out_path


def edit_image(client: OpenAI, model: str, base_path: Path, prompt: str,
               size: str, quality: str, out_path: Path) -> Path:
    logger.info("[edit] %s + %s → %s", model, base_path.name, out_path.name)
    with open(base_path, "rb") as f:
        resp = client.images.edit(
            model=model, image=f, prompt=prompt,
            size=size, quality=quality, n=1,
        )
    out_path.write_bytes(_decode_or_raise(resp, f"edit_image({out_path.name})"))
    logger.info("  saved %d KB", out_path.stat().st_size // 1024)
    return out_path


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

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Multi-step floor plan reasoning experiment (시나리오 무관)")
    p.add_argument("--project-id", required=True)
    p.add_argument("--episode-id", required=True)
    p.add_argument("--scene-indices", nargs="+", type=int, required=True)
    p.add_argument("--location-ids", nargs="+", required=True)
    p.add_argument("--target-shots", nargs="*", default=[],
                   help='format: scene:shot, e.g. "12:4 12:6 25:2"')
    p.add_argument("--text-model", default="gpt-5.5")
    p.add_argument("--image-model", default="gpt-image-2")
    p.add_argument("--image-size", default="1024x1024")
    p.add_argument("--image-quality", default="high")
    p.add_argument("--out-name", default=None,
                   help="run 디렉토리 이름 (기본: run_<timestamp>)")
    return p.parse_args()


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

    run_id = args.out_name or time.strftime("run_%Y%m%d_%H%M%S")
    out_dir = BACKEND / "scripts" / "output" / "floor_plan_test" / run_id
    out_dir.mkdir(parents=True, exist_ok=True)
    logger.info("run_id=%s out=%s", run_id, out_dir)

    client = OpenAI()

    # 0) 컨텍스트 수집
    ctx = collect_context(args.project_id, args.episode_id, args.scene_indices, args.location_ids)
    (out_dir / "context.json").write_text(json.dumps(ctx, ensure_ascii=False, indent=2), encoding="utf-8")
    logger.info("context: scenes=%d shots=%d staging=%d fixed=%d locations=%d",
                len(ctx["scenes"]), len(ctx["shots"]), len(ctx["staging"]),
                len(ctx["fixed_elements"]), len(ctx["locations"]))

    # 1) 공간 분석
    logger.info("=== Step 1: spatial analysis (%s) ===", args.text_model)
    spatial = call_llm_json(client, args.text_model, SYSTEM_STEP1, msg_step1(ctx))
    (out_dir / "step1_spatial.json").write_text(
        json.dumps(spatial, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    logger.info("space_groups=%d shared_elements=%d",
                len(spatial.get("space_groups", [])), len(spatial.get("shared_elements", [])))

    # 2) base 도면 프롬프트
    logger.info("=== Step 2: base plan prompts ===")
    plans = call_llm_json(client, args.text_model, SYSTEM_STEP2, msg_step2(ctx, spatial))
    (out_dir / "step2_base_plans.json").write_text(
        json.dumps(plans, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    logger.info("base_plans=%d", len(plans.get("base_plans", [])))

    # 2b) base 이미지 생성 — schema 검증 (HIGH-1 fix)
    base_paths: Dict[str, Path] = {}
    valid_plan_ids: set = set()
    for plan in plans.get("base_plans", []):
        pid = plan.get("id")
        prompt_str = plan.get("t2i_prompt")
        if not pid or not prompt_str:
            logger.error("base plan schema 위반 (id/t2i_prompt 누락): %s", plan)
            continue
        valid_plan_ids.add(pid)
        out_p = out_dir / f"base_{pid}.png"
        try:
            gen_image(client, args.image_model, prompt_str,
                      args.image_size, args.image_quality, out_p)
            base_paths[pid] = out_p
        except Exception as e:
            logger.error("base image failed (%s): %s", pid, e)

    # 3) 샷별 overlay
    parsed_shots = []
    for s in args.target_shots:
        try:
            sc, sh = s.split(":")
            parsed_shots.append((int(sc), int(sh)))
        except Exception:
            logger.warning("bad target-shot: %r", s)

    shot_results = []
    for scene_idx, shot_idx in parsed_shots:
        logger.info("=== Step 3: S%d_Shot%d ===", scene_idx, shot_idx)
        try:
            overlay = call_llm_json(
                client, args.text_model, SYSTEM_STEP3,
                msg_step3(ctx, spatial, plans, scene_idx, shot_idx),
            )
        except Exception as e:
            logger.error("step3 LLM failed S%d_Shot%d: %s", scene_idx, shot_idx, e)
            continue

        label = f"S{scene_idx:02d}_Shot{shot_idx}"
        (out_dir / f"step3_shot_{label}.json").write_text(
            json.dumps(overlay, ensure_ascii=False, indent=2), encoding="utf-8"
        )
        shot_results.append(overlay)

        # 이미지 — base 위에 edit. schema 검증 (HIGH-1 fix)
        base_id = overlay.get("base_plan_id")
        shot_prompt = overlay.get("t2i_prompt")
        if not base_id:
            logger.error("S%d_Shot%d: base_plan_id 누락, edit 스킵", scene_idx, shot_idx)
            continue
        if base_id not in valid_plan_ids:
            logger.error("S%d_Shot%d: base_plan_id %r가 base_plans에 없음 (valid=%s), edit 스킵",
                         scene_idx, shot_idx, base_id, sorted(valid_plan_ids))
            continue
        if not shot_prompt:
            logger.error("S%d_Shot%d: t2i_prompt 누락, edit 스킵", scene_idx, shot_idx)
            continue
        base_path = base_paths.get(base_id)
        if not base_path:
            logger.warning("S%d_Shot%d: base 이미지 생성 실패한 plan %r, edit 스킵",
                           scene_idx, shot_idx, base_id)
            continue
        out_p = out_dir / f"shot_{label}.png"
        try:
            edit_image(client, args.image_model, base_path, shot_prompt,
                       args.image_size, args.image_quality, out_p)
        except Exception as e:
            logger.error("shot edit failed S%d_Shot%d: %s", scene_idx, shot_idx, e)

    # manifest
    manifest = {
        "run_id": run_id,
        "args": vars(args),
        "models": {"text": args.text_model, "image": args.image_model},
        "stats": {
            "scenes": len(ctx["scenes"]), "shots": len(ctx["shots"]),
            "space_groups": len(spatial.get("space_groups", [])),
            "base_plans": len(plans.get("base_plans", [])),
            "shot_overlays": len(shot_results),
        },
        "files": sorted(p.name for p in out_dir.glob("*")),
    }
    (out_dir / "manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    logger.info("=== DONE === %s", out_dir)
    return 0


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