"""Spike Test 6 — LLM 자동 도면 prompt 생성 + 수동 v3 비교.

검증: Gemini Pro가 시나리오만 보고 사람 v3 수준의 정밀 도면 prompt를 자동 작성하는가?

Step 1: Gemini Pro → 도면 prompt (영문, top-down architectural)
Step 2: gpt-image-2 → 도면 PNG (자동본)
Step 3: 수동 v3 도면 PNG와 시각 비교 (sof기존 30_floorplan_v3_tworoom.png 활용)
Step 4: 두 도면 PNG 각각으로 chain bg render → 비교
"""
from __future__ import annotations

import base64
import concurrent.futures
import json
import os
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any, Dict, List, Tuple

try:
    from dotenv import load_dotenv
    load_dotenv(Path(__file__).parent.parent.parent / ".env")
except ImportError:
    pass

from openai import OpenAI


PID = "c00bbe19-a9b5-463f-acfc-806f2e820258"
EID = "fe165e3a-19c2-4a0f-9acb-e0c9bab0ee5a"
ROOT = Path("/Users/manta/Documents/Projects/TheRoad-I1")
CKPT = ROOT / "projects" / PID / "checkpoints" / "episodes" / EID
HERE = Path(__file__).parent
RESULTS = HERE / "results"
RESULTS.mkdir(parents=True, exist_ok=True)

# 비교 baseline: 수동 v3 도면 PNG (이전 실험 결과)
MANUAL_V3_FLOORPLAN = ROOT / "backend" / "scripts" / "experiment_results" / "30_floorplan_v3_tworoom.png"

ROOFTOP_SCENES = (5, 12, 14, 18, 25, 27)


# ───────── chain bg prompt (variant F와 동일) ─────────

CHAIN_BG_PROMPT = (
    "Photorealistic cinematic interior background of a Korean two-room "
    "rooftop apartment (투룸 옥탑방), pale daylight, no people, "
    "single static frame, 35mm film grain, eye-level wide angle. Match the "
    "architectural floor plan reference (top-down) EXACTLY for furniture "
    "and door positions. "
    "Camera is positioned in the LIVING/KITCHEN room near the front entry "
    "side, looking diagonally toward the back where the kitchen counter "
    "and the two interior bedroom doors are visible. "
    "Visible in this view: ONE television on a low wooden stand against "
    "the left wall (single TV — do NOT add a second TV); kitchen counter "
    "with stainless sink in the middle of the back wall; small window with "
    "thin curtain to the right of the sink; wooden dining table with two or "
    "three chairs in the center floor; the closed interior door to the main "
    "bedroom (안방) at the upper-right; the closed interior door to the "
    "daughter's bedroom (수리영 방) on the right wall further forward. "
    "DO NOT show interior of bedrooms or bathroom — only the living/kitchen "
    "room is in this frame. The apartment is one rooftop apartment, NOT a "
    "multi-floor or multi-unit building. "
    "Material/lighting: dirty plaster walls in muted yellow-gray, worn "
    "wooden floorboards, low ceiling with a single bare bulb fixture, pale "
    "daylight from the small back window, soft warm tungsten ambient, dust "
    "in the air, damp domestic shadows, uneasy stillness, muted blue-gray "
    "and washed yellow palette."
)


# ───────── Step 1: Gemini Pro 자동 prompt 생성 ─────────

PROMPT_GEN_SYSTEM = """You are a film production designer. Given the screenplay scenes that occur in a single shooting LOCATION, write a precise English text-to-image prompt that will produce a TOP-DOWN ARCHITECTURAL FLOOR PLAN diagram of that location.

The output prompt will be sent directly to gpt-image-2 (text-to-image). It must produce a schematic floor plan (NOT a photorealistic interior).

REQUIREMENTS for the prompt:
1. Start with: "Top-down architectural floor plan of [location description]."
2. List ALL rooms/zones explicitly with their relative positions (corner, side, center).
3. For each room, specify: which wall has which furniture, the size of furniture relative to the room, and how furniture relates to other furniture (e.g., "TV on left wall facing sofa on right wall").
4. Specify all doors: which wall, where it leads (interior to which room? front entry to outside what?).
5. Specify windows: which wall, which room.
6. Specify external adjacency: what is OUTSIDE the front entry door (rooftop concrete? balcony? hallway? street? cite the scenario explicitly).
7. Add explicit NEGATIVE constraints — what NOT to draw (e.g., "NOT a balcony, NOT a multi-floor building, NOT a single-room studio").
8. Use universal nouns. NO scenario proper nouns (character names, place names, work titles).
9. Korean common nouns are OK in parentheses for room labels (e.g., "main bedroom (안방)").
10. End with: "Schematic line drawing style, clean labels, no shading, white background."

CRITICAL: Read EVERY scene/shot carefully. Identify zone markers (e.g., '/거실', '/안방', '/욕실', '/현관', '/욕조'). Every zone mentioned must appear in the floor plan. If a scene shows a bathroom mirror, the floor plan must include the bathroom with mirror. If a scene shows the front door opening to a rooftop, the floor plan must show ROOFTOP outside, not balcony.

OUTPUT: One single block of English prompt text only. No JSON, no explanation, no markdown headers."""


def _build_user_prompt() -> str:
    """test_01과 동일 — 옥탑방 6 씬 + selected shots."""
    scene_save = json.loads((CKPT / "scene_save" / "manifest.json").read_text())
    sv = json.loads((CKPT / "shot_validator" / "manifest.json").read_text())["data"]
    ss = json.loads((CKPT / "shot_selection" / "manifest.json").read_text())["data"]

    segs = scene_save["data"]["segments"]
    sel_map = {s["scene_index"]: set(s.get("selected_shot_indices", []))
               for s in ss.get("scenes", [])}

    parts: List[str] = [
        "LOCATION ID: L05",
        "LOCATION DESCRIPTION (from scenario): a Korean rooftop apartment that the family lives in.",
        "",
        "── ALL SCENES OCCURRING IN THIS LOCATION ──",
        "",
    ]
    for seg in segs:
        if seg["scene_index"] in ROOFTOP_SCENES:
            parts.append(f"### Scene {seg['scene_index']} — {seg.get('heading', '')}")
            parts.append(seg.get("text", ""))
            parts.append("")

    parts.append("── SELECTED SHOTS IN THIS LOCATION (visible elements) ──")
    parts.append("")
    for s in sv.get("scenes", []):
        si = s["scene_index"]
        if si not in ROOFTOP_SCENES:
            continue
        sel = sel_map.get(si, set())
        parts.append(f"### Scene {si} selected shots:")
        for sh in s.get("shots", []):
            if sh["shot_index"] in sel:
                parts.append(f"- Shot {sh['shot_index']}: {sh.get('description','')}")
        parts.append("")

    parts.append("Now write the floor plan text-to-image prompt for this location.")
    return "\n".join(parts)


GEMINI_TEXT_API = (
    "https://generativelanguage.googleapis.com/v1beta/models/"
    "gemini-3.1-pro-preview:generateContent?key={key}"
)


def call_gemini_text(system_prompt: str, user_prompt: str) -> str:
    api_key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
    body = {
        "systemInstruction": {"parts": [{"text": system_prompt}]},
        "contents": [{"role": "user", "parts": [{"text": user_prompt}]}],
        "generationConfig": {
            "temperature": 0.3,
            "maxOutputTokens": 4096,
        },
    }
    url = GEMINI_TEXT_API.format(key=api_key)
    req = urllib.request.Request(
        url,
        data=json.dumps(body).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=600) as resp:
        payload = json.loads(resp.read().decode("utf-8"))
    return payload["candidates"][0]["content"]["parts"][0]["text"]


# ───────── Step 2: gpt-image-2 도면 PNG 생성 ─────────

def call_gpt_image_2(prompt: str) -> bytes:
    client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
    resp = client.images.generate(
        model="gpt-image-2",
        prompt=prompt,
        size="1024x1024",
        quality="high",
        n=1,
    )
    return base64.b64decode(resp.data[0].b64_json)


# ───────── Step 4: chain bg render with floor plan ref ─────────

GEMINI_IMAGE_API = (
    "https://generativelanguage.googleapis.com/v1beta/models/"
    "gemini-3.1-flash-image-preview:generateContent?key={key}"
)


def gemini_chain_bg(prompt: str, floor_plan_png: bytes) -> bytes:
    """Gemini nano-banana-2 chain bg render with floor plan ref."""
    api_key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
    parts = [
        {"text": "[architectural floor plan reference — top-down]"},
        {"inline_data": {"mime_type": "image/png",
                          "data": base64.b64encode(floor_plan_png).decode()}},
        {"text": prompt},
    ]
    body = {
        "contents": [{"parts": parts}],
        "generationConfig": {
            "responseModalities": ["TEXT", "IMAGE"],
            "imageConfig": {"aspectRatio": "16:9", "imageSize": "2K"},
        },
    }
    url = GEMINI_IMAGE_API.format(key=api_key)
    req = urllib.request.Request(
        url,
        data=json.dumps(body).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=600) as resp:
        payload = json.loads(resp.read().decode("utf-8"))
    for part in payload["candidates"][0]["content"]["parts"]:
        d = part.get("inline_data") or part.get("inlineData")
        if d:
            return base64.b64decode(d["data"])
    raise RuntimeError("no image part in chain bg response")


# ───────── Main ─────────

def main() -> int:
    # Step 1: Gemini Pro 자동 prompt 생성
    user_input = _build_user_prompt()
    print(f"[INFO] Step 1: Gemini Pro → 도면 prompt 자동 생성 (input {len(user_input)}자)...")
    auto_prompt = call_gemini_text(PROMPT_GEN_SYSTEM, user_input)
    (RESULTS / "auto_floor_plan_prompt.txt").write_text(auto_prompt, encoding="utf-8")
    print(f"  → auto_floor_plan_prompt.txt ({len(auto_prompt)}자)")
    print(f"\n--- AUTO PROMPT (preview) ---")
    print(auto_prompt[:500] + ("..." if len(auto_prompt) > 500 else ""))
    print(f"--- end ---\n")

    # Step 2: gpt-image-2 자동 도면 PNG
    print("[INFO] Step 2: gpt-image-2 → 도면 PNG (자동본 + 수동 v3 baseline)")
    auto_png_path = RESULTS / "auto_floor_plan.png"
    manual_png_bytes = MANUAL_V3_FLOORPLAN.read_bytes()

    auto_png_bytes = call_gpt_image_2(auto_prompt)
    auto_png_path.write_bytes(auto_png_bytes)
    print(f"  자동 도면 → {auto_png_path} ({len(auto_png_bytes):,} bytes)")
    print(f"  수동 v3 도면(baseline) ← {MANUAL_V3_FLOORPLAN} ({len(manual_png_bytes):,} bytes)")

    # Step 4: chain bg render with each floor plan (병렬)
    print("\n[INFO] Step 4: 두 도면 각각으로 chain bg render (Gemini nano-banana-2)...")
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as ex:
        f_auto = ex.submit(gemini_chain_bg, CHAIN_BG_PROMPT, auto_png_bytes)
        f_manual = ex.submit(gemini_chain_bg, CHAIN_BG_PROMPT, manual_png_bytes)
        results: Dict[str, Tuple[str, bytes | str]] = {}
        for name, fut in (("auto", f_auto), ("manual_v3", f_manual)):
            try:
                results[name] = ("ok", fut.result(timeout=600))
            except Exception as exc:
                results[name] = ("error", f"{type(exc).__name__}: {exc}")

    for name, (status, payload) in results.items():
        if status == "ok":
            out = RESULTS / f"chain_bg_{name}.png"
            out.write_bytes(payload)
            print(f"  [OK] chain_bg_{name} → {out} ({len(payload):,} bytes)")
        else:
            print(f"  [ERROR] chain_bg_{name}: {payload}")

    return 0


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