"""S5_Shot3 chain bg 재생성 — floor plan 기반 (variant D).

Step 1: floor plan 도면 + photoreal description으로 chain bg PNG 재생성.
        gpt-image-2 images.edit (도면을 reference로) 호출.
        목표: 단칸 옥탑방 (multi-room 아님), 도면 layout 정확히 따름, 1 TV.

Step 2: 새 chain bg + 도면 + C04 + P06 = 4 ref + 기존 prompt(A 그대로)로
        nano-banana-2 호출 → variant D.

가설: chain bg PNG 자체가 도면 기반으로 정확히 그려지면,
  - prompt의 "TV 재정의"도 chain bg에 이미 있는 TV 한 개로 통합 (이중 합성 사라짐)
  - 옥탑방이 multi-room 빌라로 변형되는 근본 원인 제거 (S25_Shot13 케이스도 fix)
"""
from __future__ import annotations

import base64
import json
import os
import socket
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

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

from openai import OpenAI


PID = "c00bbe19-a9b5-463f-acfc-806f2e820258"
EID = "fe165e3a-19c2-4a0f-9acb-e0c9bab0ee5a"
PROJECT_ROOT = Path("/Users/manta/Documents/Projects/TheRoad-I1")

# 기존 references (재사용)
C04_FACE_REF = (
    PROJECT_ROOT / "projects" / PID / "images" / EID
    / "reference" / "008b5d4c-5a02-4233-853e-ffc99225cfe1.png"
)
P06_REF = (
    PROJECT_ROOT / "projects" / PID / "images" / EID
    / "reference" / "48735792-33c0-415b-afdf-46146fe19ed4.png"
)

OUT_DIR = Path(__file__).parent / "experiment_results"
FLOORPLAN_PATH = OUT_DIR / "00_floorplan.png"  # 이전 실험에서 생성됨

# Step 1: floor plan 기반 chain bg 재생성 prompt
# 도면 layout을 photoreal interior로 변환. 핵심: "single rooftop room", "low ceiling",
# "front door opens directly to outdoor rooftop", "TV in one corner only" 명시.
NEW_CHAIN_BG_PROMPT = (
    "Photorealistic cinematic interior background. Convert the architectural "
    "floor plan reference (top-down diagram) into a 35mm film-style photoreal "
    "view of the actual room from inside, eye-level wide angle, single static "
    "frame, no people. Match the reference floor plan layout EXACTLY: "
    "this is ONE single rooftop room (옥탑방) on top of a building — NOT a "
    "multi-room apartment, NOT a living-bedroom-kitchen suite. The whole "
    "interior is a single 5m-wide room with the layout described below. "
    "Layout from this view (camera near the front-right corner looking toward "
    "the back-left): "
    "Left wall: ONE small old CRT television on a low wooden stand (single TV, "
    "do not draw a second TV anywhere) and a worn fabric sofa in front of it. "
    "Back wall: kitchen counter with a stainless sink in the middle, a small "
    "window with a thin curtain to the right of the sink, and a narrow plain "
    "interior doorway (closed) at the far back-right corner — this doorway "
    "leads to a tiny adjoining bedroom space (do NOT show its inside). "
    "Right wall: a simple wooden dining table with two or three chairs. "
    "Front-right (near the camera): the outer entry door is offscreen-right, "
    "ignore it for this view. "
    "Material/lighting: dirty plaster walls in muted yellow-gray, worn wooden "
    "floorboards, low ceiling with a single bare bulb fixture, pale daylight "
    "leaking through the small back window, soft warm tungsten ambient light "
    "filling the rest, dust in the air, damp domestic shadows, uneasy "
    "stillness, 35mm film grain, muted blue-gray and washed yellow palette."
)

# Step 2: 기존 prompt (A 그대로 — 안내 X)
PROMPT_A_ORIGINAL = (
    "Photorealistic cinematic still. "
    "[L05: A cramped modern Korean 옥탑방 (rooftop room) interior, dusty small "
    "window, worn sink, low table, pale daylight, damp muted domestic shadows.] "
    "C04O06 in a plain apron occupies the right-center of the frame, at "
    "three-quarter angle from the left, one hand gripping the apron hem, eyes "
    "fixed on the television at the far left. The lens emphasizes her serious "
    "expression, tightened jaw, and the tense hand at the apron hem through a "
    "thin veil of white kitchen steam in the left foreground. A blurred "
    "television edge, screen angled away, sits at the far left, its cold "
    "blue-gray glow brushing her cheek; pale daylight from the closed dusty "
    "window adds washed yellow highlights. A steaming pot with a slightly "
    "lifted lid sits in the foreground, P06 rests rinsed beside the sink, and "
    "a thin stream from the faucet falls into the basin in the background. "
    "Muted blue-gray and washed yellow palette, damp domestic surfaces, uneasy "
    "stillness."
)

# Gemini call
GEMINI_API_URL_TEMPLATE = (
    "https://generativelanguage.googleapis.com/v1beta/models/{model}:"
    "generateContent?key={api_key}"
)
GEMINI_MODEL = "gemini-3.1-flash-image-preview"
ASPECT_RATIO = "16:9"
TIMEOUT = 240
MAX_RETRIES = 2


def gemini_generate(api_key: str, prompt: str, input_images: list) -> bytes:
    parts: list = [{"text": prompt}]
    for label, img_bytes in input_images:
        if label:
            parts.append({"text": label})
        parts.append({
            "inline_data": {
                "mime_type": "image/png",
                "data": base64.b64encode(img_bytes).decode("ascii"),
            }
        })
    body = {
        "contents": [{"parts": parts}],
        "generationConfig": {
            "responseModalities": ["TEXT", "IMAGE"],
            "imageConfig": {"aspectRatio": ASPECT_RATIO, "imageSize": "2K"},
        },
    }
    url = GEMINI_API_URL_TEMPLATE.format(model=GEMINI_MODEL, api_key=api_key)
    req = urllib.request.Request(
        url, data=json.dumps(body).encode("utf-8"),
        headers={"Content-Type": "application/json"}, method="POST",
    )
    last = None
    for attempt in range(1, MAX_RETRIES + 2):
        try:
            with urllib.request.urlopen(req, timeout=TIMEOUT) as resp:
                payload = json.loads(resp.read().decode("utf-8"))
            break
        except urllib.error.HTTPError as exc:
            text = exc.read().decode("utf-8", errors="replace")
            last = RuntimeError(f"HTTP {exc.code}: {text}")
            if exc.code in {429, 500, 502, 503, 504} and attempt <= MAX_RETRIES:
                time.sleep(2 * attempt); continue
            raise last from exc
        except (urllib.error.URLError, socket.timeout) as exc:
            last = exc
            if attempt <= MAX_RETRIES:
                time.sleep(2 * attempt); continue
            raise RuntimeError(f"URL error: {exc}") from exc
    else:
        raise RuntimeError(f"Failed: {last}")
    pf = payload.get("promptFeedback", {})
    if pf.get("blockReason"):
        raise RuntimeError(f"Moderation: {pf['blockReason']}")
    cands = payload.get("candidates", [])
    if cands and cands[0].get("finishReason") == "SAFETY":
        raise RuntimeError("Safety filter")
    for cand in cands:
        for part in cand.get("content", {}).get("parts", []):
            inline = part.get("inlineData") or part.get("inline_data")
            if isinstance(inline, dict) and inline.get("data"):
                return base64.b64decode(inline["data"])
    raise RuntimeError(f"No image: {json.dumps(payload)[:300]}")


def main():
    openai_key = os.environ.get("OPENAI_API_KEY", "")
    gemini_key = os.environ.get("GEMINI_API_KEY", "")
    if not openai_key or not gemini_key:
        print("ERROR: OPENAI_API_KEY + GEMINI_API_KEY required", file=sys.stderr)
        sys.exit(1)

    if not FLOORPLAN_PATH.exists():
        print(f"ERROR: floor plan not found at {FLOORPLAN_PATH}", file=sys.stderr)
        print("Run experiment_chain_bg_with_floorplan.py first to generate it.")
        sys.exit(1)

    OUT_DIR.mkdir(exist_ok=True)

    # ── Step 1: gpt-image-2로 도면 기반 chain bg 재생성 ──
    print("=" * 60)
    print("Step 1: chain bg 재생성 (gpt-image-2 + 도면 ref)")
    print("=" * 60)
    print(f"prompt length: {len(NEW_CHAIN_BG_PROMPT)} chars")
    new_chain_bg_path = OUT_DIR / "10_new_chain_bg_from_floorplan.png"
    if new_chain_bg_path.exists():
        print(f"이미 존재 — 재사용 (지우려면 파일 삭제): {new_chain_bg_path}")
    else:
        client = OpenAI(api_key=openai_key)
        with open(FLOORPLAN_PATH, "rb") as f:
            resp = client.images.edit(
                model="gpt-image-2",
                image=f,
                prompt=NEW_CHAIN_BG_PROMPT,
                size="1536x1024",
                quality="high",
                n=1,
            )
        b64 = resp.data[0].b64_json
        new_chain_bg_path.write_bytes(base64.b64decode(b64))
        print(f"OK: {new_chain_bg_path} ({new_chain_bg_path.stat().st_size:,} bytes)")
    print()

    # ── Step 2: variant D — nano-banana-2 + 새 chain bg + 도면 + C04 + P06 ──
    print("=" * 60)
    print("Step 2: variant D — 새 chain bg + 도면 + 인물/소품 + 기존 prompt(A)")
    print("=" * 60)
    input_images = [
        (
            "spatial layout reference (floor plan, top-down) — "
            "use this to understand the room layout",
            FLOORPLAN_PATH.read_bytes(),
        ),
        (
            "background reference (rebuilt from floor plan) — "
            "match wall/floor/ceiling/lighting and ALL furniture positions; "
            "DO NOT add or relocate any furniture",
            new_chain_bg_path.read_bytes(),
        ),
        ("character C04 identity", C04_FACE_REF.read_bytes()),
        ("object P06", P06_REF.read_bytes()),
    ]
    print(f"References (4):")
    for i, (label, _) in enumerate(input_images, 1):
        print(f"  [{i}] {label[:60]}...")
    print(f"prompt length: {len(PROMPT_A_ORIGINAL)} chars (A 그대로 — 안내 X)")
    print()

    d_path = OUT_DIR / "s5shot3_D_floorplan_rebuilt_chainbg.png"
    try:
        d_bytes = gemini_generate(gemini_key, PROMPT_A_ORIGINAL, input_images)
        d_path.write_bytes(d_bytes)
        print(f"OK: {d_path} ({len(d_bytes):,} bytes)")
    except Exception as exc:
        print(f"FAIL: {exc}")
        d_path = None
    print()

    meta = {
        "test": "S5_Shot3 — chain bg rebuilt from floor plan (variant D)",
        "step1_chain_bg_rebuild": {
            "model": "gpt-image-2",
            "input_ref": str(FLOORPLAN_PATH),
            "prompt": NEW_CHAIN_BG_PROMPT,
            "result_path": str(new_chain_bg_path),
        },
        "step2_variant_D": {
            "model": GEMINI_MODEL,
            "references": [
                "00_floorplan.png",
                "10_new_chain_bg_from_floorplan.png",
                "C04 face",
                "P06 prop",
            ],
            "prompt": PROMPT_A_ORIGINAL,
            "added_section": "(none — A 그대로)",
            "result_path": str(d_path) if d_path else None,
        },
    }
    (OUT_DIR / "s5shot3_floorplan_rebuild_metadata.json").write_text(
        json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    print(f"metadata: {OUT_DIR / 's5shot3_floorplan_rebuild_metadata.json'}")


if __name__ == "__main__":
    main()