"""실험: 옥탑방 도면 base 2×2 매트릭스 비교

조합 4:
- GPT-5.4 분석 → gpt-image-2 생성
- GPT-5.4 분석 → Gemini 3.1 Flash Image 생성
- Gemini Pro 분석 → gpt-image-2 생성
- Gemini Pro 분석 → Gemini 3.1 Flash Image 생성

스케일/공간 교정 시스템 프롬프트:
- 옥탑방 10~12평 소규모
- 거실·주방 통합 + 민숙 방 + 수리영 방 (벽+문으로 분리)
- 수리영 방 **내부** 침대 뒤 **커튼 알코브** (시체 숨은 위치)
- 원문 명시 소품만 (TV+스탠드, 소파, 식탁+찻잔 2개, 싱크·스토브+냄비, 침대 2개, 커튼, 백팩)
- base는 인물 없음, 도면 스타일 (top-down, line drawing)

출력: backend/scripts/output/floor_plan_compare/
  - prompts_gpt.json / prompts_gemini.json
  - base_gpt__openai.png
  - base_gpt__gemini.png
  - base_gemini__openai.png
  - base_gemini__gemini.png
"""
from __future__ import annotations

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

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
from app.modules.llm.gemini_text_client import GeminiTextClient  # noqa: E402
from app.modules.llm.gemini_image_client import GeminiImageClient  # noqa: E402

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

PROJECT_ID = "d8a3b254-45c9-4eb4-bba6-8bae95eca96f"
EPISODE_ID = "e71c5a18-88e7-40e2-81bd-dc28c56a62f7"
OUT_DIR = BACKEND / "scripts" / "output" / "floor_plan_compare"
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"]

# ─ 모델
OPENAI_TEXT = "gpt-5.5"  # 2026-04-23 출시, 1M context
GEMINI_TEXT = "gemini-3.1-pro-preview"  # gemini_text_client 기본
IMAGE_OPENAI = "gpt-image-2"
IMAGE_SIZE = "1024x1024"
IMAGE_QUALITY = "high"

openai_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:
        return json.loads(arches[0].read_text(encoding="utf-8"))
    raise FileNotFoundError(step_id)


def collect_context() -> Dict[str, Any]:
    ctx: Dict[str, Any] = {"scope": {"scenes": OKTAP_SCENES, "locations": OKTAP_LOCATIONS}}

    scene_save = load_cp("scene_save")
    ctx["scenes"] = [
        {"scene_index": s["scene_index"], "heading": s["heading"], "text": s["text"]}
        for s in scene_save["data"]["segments"]
        if s["scene_index"] in OKTAP_SCENES
    ]

    loc_cp = load_cp("entity_extract_location")
    locs = {l["short_id"]: l for l in loc_cp["data"]["locations"]}
    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 OKTAP_LOCATIONS if lid in locs
    ]

    sc_cp = load_cp("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_cp["data"]["scenes"]
        if s["scene_index"] in OKTAP_SCENES
    ]

    return ctx


# ──────────────────────────────────────────────────────────
# 시스템 프롬프트 — 스케일/커튼/원문 소품 교정
# ──────────────────────────────────────────────────────────

SYSTEM_PROMPT = """당신은 건축 도면 전문가입니다. 한국 드라마 시나리오 장면을 **건축 도면(architectural floor plan) T2I 프롬프트**로 변환합니다.

## 공간 구조 강제 규칙 (시나리오 분석 결과, 반드시 반영)
- **옥탑방 = 10~12평 규모의 작은 한국형 주거 공간**. 과도하게 크게 그리지 말 것.
- **레이아웃**:
  - 현관(entry) — 입구
  - 거실·주방 통합 — 식탁(찻잔 2개, 하나 엎어짐)+의자 2개 / 싱크+스토브(냄비)/ 소파 + TV (낮은 스탠드) + 플로어 스탠드
  - **민숙 방** — 거실에서 문으로 진입하는 **작은** 침실 (침대 하나)
  - **수리영 방** — 거실에서 문으로 진입하는 **작은** 침실 (침대 하나)
  - **수리영 방 내부** 침대 뒤에 **커튼(fabric partition)으로 구분된 작은 알코브** — 시체 발견 위치, 매우 중요
  - 좁은 화장실 정도(optional)
- **시나리오 명시 소품만 그리기** (추측 금지): TV, 거실 스탠드, 소파(optional), 식탁+찻잔 2개, 싱크+스토브+냄비, 침대×2, 커튼 알코브, 백팩+인형(수리영 방 바닥)
- **옥상/외곽/계단/지상은 base에 포함하지 않음** (별도 도면 대상)

## 도면 스타일
- 흑백 기술 도면, top-down orthographic view, 벽=검정 실선, 문=swing arc, 창=이중 평행선
- 가구 단순 외곽선 + 한영 라벨 (예: "Living · 거실", "Minsook Room · 민숙 방", "Curtain Alcove · 커튼 알코브")
- 우측 하단에 compass(N) + scale bar(0-5m)
- base는 **인물 없음** (empty of characters)

## 출력 (JSON, 엄격)
{"base_prompt": "string (영문, 위 레이아웃/스타일 지시 전부 포함)"}

추가 설명 금지.
"""


def build_user_message(ctx: Dict[str, Any]) -> str:
    return (
        "## 옥탑방 관련 씬 원문 + fixed_elements + location description\n\n"
        f"### 씬 원문 ({len(ctx['scenes'])}개, 자르지 않음)\n"
        + "\n\n".join(f"【{s['heading']}】\n{s['text']}" for s in ctx["scenes"])
        + "\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"  - [{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"
        "위 정보를 바탕으로 **base_prompt JSON**을 출력하세요. 시스템의 공간 구조 강제 규칙을 반드시 따를 것 (옥탑방 작게, 수리영 방 안 커튼 알코브)."
    )


# ──────────────────────────────────────────────────────────
# 분석 (텍스트)
# ──────────────────────────────────────────────────────────

def analyze_gpt(ctx: Dict[str, Any]) -> Dict[str, Any]:
    logger.info("[analyze] GPT-5.4...")
    resp = openai_client.chat.completions.create(
        model=OPENAI_TEXT,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": build_user_message(ctx)},
        ],
        response_format={"type": "json_object"},
    )
    data = json.loads(resp.choices[0].message.content)
    logger.info("[analyze] GPT-5.4 base_prompt 길이=%d", len(data.get("base_prompt", "")))
    return data


def analyze_gemini(ctx: Dict[str, Any]) -> Dict[str, Any]:
    logger.info("[analyze] Gemini Pro...")
    client = GeminiTextClient()
    schema = {
        "type": "object",
        "properties": {
            "base_prompt": {"type": "string"}
        },
        "required": ["base_prompt"],
    }
    data = client.send_structured(
        user_message=build_user_message(ctx),
        response_schema=schema,
        schema_name="FloorPlanBase",
        system_instruction=SYSTEM_PROMPT,
    )
    logger.info("[analyze] Gemini base_prompt 길이=%d", len(data.get("base_prompt", "")))
    return data


# ──────────────────────────────────────────────────────────
# 이미지 생성
# ──────────────────────────────────────────────────────────

def gen_openai(prompt: str, out_path: Path) -> Path:
    logger.info("[image] openai gpt-image-2 → %s", out_path.name)
    resp = openai_client.images.generate(
        model=IMAGE_OPENAI, prompt=prompt, size=IMAGE_SIZE, quality=IMAGE_QUALITY, n=1,
    )
    out_path.write_bytes(base64.b64decode(resp.data[0].b64_json))
    logger.info("  saved %d KB", out_path.stat().st_size // 1024)
    return out_path


def gen_gemini(prompt: str, out_path: Path) -> Path:
    logger.info("[image] gemini 3.1 flash image → %s", out_path.name)
    client = GeminiImageClient()
    client.set_context(project_id=PROJECT_ID, episode_id=EPISODE_ID,
                       step_name="experiment_floor_plan_compare")
    img_bytes, ms = client.generate_image(prompt, aspect_ratio="1:1")
    out_path.write_bytes(img_bytes)
    logger.info("  saved %d KB (%d ms)", out_path.stat().st_size // 1024, ms)
    return out_path


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

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

    ctx = collect_context()
    (OUT_DIR / "context.json").write_text(
        json.dumps(ctx, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    logger.info("context saved (scenes=%d, fixed=%d, locations=%d)",
                len(ctx["scenes"]), len(ctx["fixed_elements"]), len(ctx["locations"]))

    # GPT-5.5 분석만 (Gemini 제외)
    prompts_gpt = analyze_gpt(ctx)
    (OUT_DIR / "prompts_gpt.json").write_text(
        json.dumps(prompts_gpt, ensure_ascii=False, indent=2), encoding="utf-8"
    )

    # gpt-image-2 1장 (Gemini 이미지 제외)
    base_prompt = prompts_gpt["base_prompt"]
    gen_openai(base_prompt, OUT_DIR / "base_gpt55__openai.png")

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


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