"""S5_Shot3 chain bg 가설 검증 — floor plan 도면 추가 (variant C).

Step 1: gpt-image-2로 chain bg PNG의 layout을 표현하는 top-down floor plan 도면 생성.
Step 2: 같은 references 3장 + floor plan 1장 = 총 4장 ref + 기존 prompt(A 그대로)로
        nano-banana-2 호출 → variant C.

가설: 도면이 강한 spatial anchor가 되어 t2i_prompt의 "TV 재정의"보다 우선시됨.
      → A보다 chain bg layout을 더 충실히 재현 (B 안내와 같은 효과 또는 더 강한 효과)

비교 결과:
  A. 3 ref + 기존 prompt (안내 X)  — TV 거대 클로즈업
  B. 3 ref + [배경 안내] 추가      — chain bg 보존
  C. 4 ref (도면 추가) + 기존 prompt (안내 X) — 도면이 anchor 역할?
"""
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")

CHAIN_BG = (
    PROJECT_ROOT / "projects" / PID / "images" / EID
    / "background_chain" / "L05" / "interior_living_kitchen_day_normal.png"
)
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"
)

# Step 1: floor plan 도면 생성 prompt
# chain bg PNG의 layout을 top-down 도면으로 표현 (TV 왼쪽, 싱크대 뒤, 식탁 오른쪽, 창 가운데 뒤, 문 오른쪽 등).
FLOOR_PLAN_PROMPT = (
    "Top-down architectural floor plan diagram, hand-drawn black ink lines on "
    "off-white paper, single small Korean apartment interior layout viewed "
    "directly from above. Layout: roughly square living-kitchen-dining room "
    "about 5 meters wide. "
    "Left wall: a small CRT television on a low stand placed against the wall, "
    "labeled 'TV', a fabric sofa parallel to the wall facing the TV labeled "
    "'SOFA'. "
    "Back wall (upper edge of the plan): kitchen counter with sink in the "
    "middle labeled 'SINK', a small window with curtain immediately to the "
    "right of the sink labeled 'WINDOW', and a doorway opening to a side room "
    "at the upper-right corner labeled 'BEDROOM DOOR'. "
    "Right wall: dining table with four chairs labeled 'TABLE'. "
    "Front-right (lower-right corner): front entry door swinging outward, "
    "labeled 'FRONT DOOR'. "
    "Floor surface: wooden floorboards indicated by parallel light line "
    "hatching. "
    "Style: clean schematic architectural diagram, plain black ink linework "
    "on warm off-white background, English labels in small caps, no "
    "perspective, no shading, no people."
)

# 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 = 180
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)

    out_dir = Path(__file__).parent / "experiment_results"
    out_dir.mkdir(exist_ok=True)

    # ── Step 1: gpt-image-2로 floor plan 도면 생성 ──
    print("=" * 60)
    print("Step 1: floor plan 도면 생성 (gpt-image-2 text-to-image)")
    print("=" * 60)
    floorplan_path = out_dir / "00_floorplan.png"
    if floorplan_path.exists():
        print(f"이미 존재 — 재사용: {floorplan_path}")
    else:
        client = OpenAI(api_key=openai_key)
        print(f"prompt length: {len(FLOOR_PLAN_PROMPT)} chars")
        resp = client.images.generate(
            model="gpt-image-2",
            prompt=FLOOR_PLAN_PROMPT,
            size="1024x1024",
            quality="high",
            n=1,
        )
        b64 = resp.data[0].b64_json
        floorplan_path.write_bytes(base64.b64decode(b64))
        print(f"OK: {floorplan_path} ({floorplan_path.stat().st_size:,} bytes)")
    print()

    # ── Step 2: variant C — nano-banana-2 + 4 ref + 기존 prompt ──
    print("=" * 60)
    print("Step 2: variant C — chain bg + C04 + P06 + floor plan + 기존 prompt(A)")
    print("=" * 60)
    print(f"Model: {GEMINI_MODEL}")
    input_images = [
        # 순서: floor plan을 먼저 두면 spatial anchor로 강조될 가능성
        (
            "spatial layout reference (floor plan, top-down) — "
            "use this to understand the room layout and furniture positions",
            floorplan_path.read_bytes(),
        ),
        (
            "background chain ref (interior_living_kitchen_day_normal for L05) — "
            "match wall/floor/ceiling/lighting",
            CHAIN_BG.read_bytes(),
        ),
        ("character C04 identity", C04_FACE_REF.read_bytes()),
        ("object P06", P06_REF.read_bytes()),
    ]
    print("References (4):")
    for i, (label, img) in enumerate(input_images, 1):
        print(f"  [{i}] {label[:60]}... ({len(img):,} bytes)")
    print()

    c_path = out_dir / "s5shot3_C_with_floorplan.png"
    print(f"prompt length: {len(PROMPT_A_ORIGINAL)} chars (A 그대로 — 안내 X)")
    try:
        c_bytes = gemini_generate(gemini_key, PROMPT_A_ORIGINAL, input_images)
        c_path.write_bytes(c_bytes)
        print(f"OK: {c_path} ({len(c_bytes):,} bytes)")
    except Exception as exc:
        print(f"FAIL: {exc}")
        c_path = None
    print()

    meta = {
        "test": "S5_Shot3 floor plan as anchor (variant C)",
        "step1_floor_plan": {
            "model": "gpt-image-2",
            "prompt": FLOOR_PLAN_PROMPT,
            "result_path": str(floorplan_path),
        },
        "step2_variant_C": {
            "model": GEMINI_MODEL,
            "references": [
                "00_floorplan.png (top-down)",
                "01_chain_bg.png (existing)",
                "C04 face ref",
                "P06 prop ref",
            ],
            "prompt": PROMPT_A_ORIGINAL,
            "added_section": "(none — A 그대로, 안내 X)",
            "result_path": str(c_path) if c_path else None,
        },
    }
    meta_path = out_dir / "s5shot3_floorplan_metadata.json"
    meta_path.write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"metadata: {meta_path}")


if __name__ == "__main__":
    main()