"""Spike Test 3 (Variant D) — 자연어 spatial DSL in prompt.

가설: 자연어 spatial language("LEFT wall", "FACES across", "upper-right corner")
      가 SVG XML보다 모델 친화적. 학습 데이터에 풍부.

대상: chain bg 노드 interior_living_kitchen_day_normal (variant F와 동일 비교).
모델: gpt-image-2 + Gemini nano-banana-2 (text-to-image, no ref).
"""
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


HERE = Path(__file__).parent
RESULTS = HERE / "results"
RESULTS.mkdir(parents=True, exist_ok=True)


# ───────── 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. "
    "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."
)


# ───────── spec.json → 자연어 spatial DSL ─────────

def _bbox(polygon: List[List[float]]) -> Tuple[float, float, float, float]:
    """polygon → (xmin, ymin, xmax, ymax)"""
    xs = [p[0] for p in polygon]
    ys = [p[1] for p in polygon]
    return min(xs), min(ys), max(xs), max(ys)


def _h_pos(cx: float, cw: float) -> str:
    if cx < cw / 3:
        return "left"
    if cx > 2 * cw / 3:
        return "right"
    return "center"


def _v_pos(cy: float, ch: float) -> str:
    if cy < ch / 3:
        return "lower (front, near front entry)"
    if cy > 2 * ch / 3:
        return "upper (back)"
    return "middle"


def _wall_label(wall: str) -> str:
    """wall name → 자연어 + camera-relative hint"""
    wmap = {
        "left": "LEFT wall (west wall, on the camera-left in this view)",
        "right": "RIGHT wall (east wall, on the camera-right in this view)",
        "top": "BACK wall (north wall, far from camera)",
        "bottom": "FRONT wall (south wall, behind camera)",
        "interior": "interior partition wall",
        "": "wall (unspecified)",
    }
    return wmap.get(wall, wall.upper() + " wall")


def _wall_short(wall: str) -> str:
    return {
        "left": "LEFT",
        "right": "RIGHT",
        "top": "BACK",
        "bottom": "FRONT",
        "interior": "interior",
    }.get(wall, wall.upper())


def spec_to_natural_dsl(spec: Dict[str, Any]) -> str:
    """spec.json → 자연어 spatial DSL.

    원칙:
    - 모델이 학습한 spatial language 사용 ("LEFT wall", "FACES the TV across")
    - 좌표는 보조 (구체 m 단위 명시)
    - top-down view 명시 (+x east, +y north)
    - rejected layouts → "do NOT generate" 명시
    """
    canvas = spec.get("canvas", {"width": 10, "height": 10})
    cw = float(canvas["width"])
    ch = float(canvas["height"])

    parts: List[str] = []
    parts.append(
        f"APARTMENT FLOOR LAYOUT (top-down view, {cw:.0f}m wide × {ch:.0f}m deep). "
        f"In this top-down system: +x = east (right), +y = north (back). "
        f"'LEFT wall' = west, 'RIGHT wall' = east, 'BACK wall' = north (far), "
        f"'FRONT wall' = south (near front entry door)."
    )
    parts.append("")

    # ROOMS
    parts.append("ROOMS (with positions, sizes, and furniture):")
    rooms_by_id = {r["id"]: r for r in spec.get("rooms", [])}
    for room in spec.get("rooms", []):
        rid = room["id"]
        xmin, ymin, xmax, ymax = _bbox(room["polygon"])
        rw = xmax - xmin
        rh = ymax - ymin
        rcx = (xmin + xmax) / 2
        rcy = (ymin + ymax) / 2

        v_label = "lower" if rcy < ch / 3 else ("upper" if rcy > 2 * ch / 3 else "middle")
        h_label = "left" if rcx < cw / 3 else ("right" if rcx > 2 * cw / 3 else "center")
        if h_label == "center" and v_label == "middle":
            pos_label = "CENTER of apartment"
        elif v_label == "middle":
            pos_label = f"{h_label.upper()} side"
        elif h_label == "center":
            pos_label = f"{v_label.upper()} center"
        else:
            pos_label = f"{v_label.upper()}-{h_label.upper()} corner"

        label_en = room.get("label_en", rid)
        label_ko = room.get("label_ko", "")

        parts.append(f"")
        parts.append(f"### {label_en.upper()} ({label_ko}) [{rid}]")
        parts.append(f"  Position: {pos_label}.")
        parts.append(f"  Size: {rw:.1f}m wide (E-W) × {rh:.1f}m deep (N-S).")

        # furniture in this room
        room_furn = [f for f in spec.get("furniture", []) if f.get("room") == rid]
        if not room_furn:
            parts.append("  Furniture: (none specified)")
            continue

        parts.append("  Furniture:")
        for f in room_furn:
            ftype = f.get("type", "item")
            label = f.get("label_en", ftype)
            wall = f.get("wall", "")
            sw, sh = f.get("size", [0.5, 0.5])
            fx, fy = f.get("position", [0, 0])

            # 가구가 어느 wall에 붙어있는지 자연어
            wall_desc = _wall_label(wall) if wall else ""
            # 가구의 room 내 상대 위치 (가운데/끝)
            rel_x = (fx - xmin) / max(rw, 0.01)
            rel_y = (fy - ymin) / max(rh, 0.01)
            if 0.4 < rel_x < 0.6 and 0.4 < rel_y < 0.6:
                rel_pos = "center of the room"
            elif wall:
                # wall이 있으면 그 wall 위 어디인지
                if wall in ("left", "right"):
                    if rel_y < 0.35:
                        rel_pos = "near the front of the wall (closer to front entry)"
                    elif rel_y > 0.65:
                        rel_pos = "near the back of the wall"
                    else:
                        rel_pos = "centered along the wall"
                elif wall in ("top", "bottom"):
                    if rel_x < 0.35:
                        rel_pos = "near the LEFT end of the wall"
                    elif rel_x > 0.65:
                        rel_pos = "near the RIGHT end of the wall"
                    else:
                        rel_pos = "centered along the wall"
                else:
                    rel_pos = "on the wall"
            else:
                rel_pos = f"at room-relative ({rel_x:.0%}, {rel_y:.0%})"

            line = (
                f"    - {label.upper()} ({ftype}): "
                f"{sw:.1f}m × {sh:.1f}m"
            )
            if wall_desc:
                line += f", on the {wall_desc}, {rel_pos}"
            else:
                line += f", {rel_pos}"
            line += "."
            parts.append(line)

    # DOORS
    parts.append("")
    parts.append("DOORS:")
    for door in spec.get("doors", []):
        kind = door.get("kind", "interior")
        from_r = door.get("from_room", "")
        to_r = door.get("to_room", "")
        from_label = rooms_by_id.get(from_r, {}).get("label_en", from_r) or "exterior"
        to_label = rooms_by_id.get(to_r, {}).get("label_en", to_r) or "exterior"
        if kind == "front_entry":
            parts.append(
                f"  - FRONT ENTRY DOOR: connects {from_label} to OUTSIDE ({to_label}). "
                f"This is the only door leaving the apartment."
            )
        elif kind == "window":
            parts.append(f"  - WINDOW in {from_label}.")
        else:
            parts.append(f"  - INTERIOR DOOR: connects {from_label} ↔ {to_label}.")

    # EXTERIOR
    if spec.get("exterior_adjacency_zones"):
        parts.append("")
        parts.append("OUTSIDE THE APARTMENT (what is beyond the front entry door):")
        for ez in spec["exterior_adjacency_zones"]:
            parts.append(
                f"  - {ez.get('label_en','').upper()} ({ez.get('type','')}): "
                f"this is what is visible through the front entry door. "
                f"NOT a balcony, NOT a hallway, NOT a street."
            )

    # CAMERA-RELATIVE FACING
    parts.append("")
    parts.append(
        "FURNITURE FACING RULES (important for photoreal rendering):"
    )
    parts.append(
        "  - The TV (on the LEFT wall) and the SOFA (if any, on the RIGHT wall) "
        "FACE EACH OTHER ACROSS the room. The dining table sits between them in the center."
    )
    parts.append(
        "  - The kitchen SINK on the BACK wall faces toward the front (toward the camera/front entry)."
    )
    parts.append(
        "  - The window on the BACK wall is a small punched opening, not a wall of glass."
    )

    # REJECTED LAYOUTS
    if spec.get("rejected_layouts"):
        parts.append("")
        parts.append("FORBIDDEN LAYOUTS (do NOT generate any of these):")
        for r in spec["rejected_layouts"]:
            parts.append(f"  - {r}")

    return "\n".join(parts)


# ───────── Build variant D prompt ─────────

def build_variant_d_prompt(dsl_text: str) -> str:
    return (
        "[ARCHITECTURAL FLOOR LAYOUT — natural-language spatial description, "
        "use as authoritative spatial layout]\n"
        f"{dsl_text}\n\n"
        "[RENDER INSTRUCTIONS]\n"
        + CHAIN_BG_PROMPT
        + "\n\nMatch the architectural layout above EXACTLY for furniture "
        "positions, door positions, and window positions. Do not invent "
        "additional furniture or doors. The TV must be on the LEFT wall "
        "(camera-left). The sink on the BACK wall (far from camera). The "
        "two interior bedroom doors on the RIGHT wall and at the upper-right. "
        "Render as a single photorealistic eye-level cinematic frame from "
        "the camera position described in the prompt above (NOT a top-down view)."
    )


# ───────── gpt-image-2 ─────────

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="1536x1024",
        quality="high",
        n=1,
    )
    return base64.b64decode(resp.data[0].b64_json)


# ───────── Gemini nano-banana-2 ─────────

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


def call_gemini_nano_banana(prompt: str) -> bytes:
    api_key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
    body = {
        "contents": [{"parts": [{"text": prompt}]}],
        "generationConfig": {
            "responseModalities": ["TEXT", "IMAGE"],
            "imageConfig": {"aspectRatio": "16:9", "imageSize": "2K"},
        },
    }
    url = GEMINI_API_URL.format(key=api_key)
    last_err: Exception | None = None
    for attempt in range(1, 4):
        try:
            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"]:
                if "inline_data" in part:
                    return base64.b64decode(part["inline_data"]["data"])
                if "inlineData" in part:
                    return base64.b64decode(part["inlineData"]["data"])
            raise RuntimeError("no image part in response")
        except urllib.error.HTTPError as exc:
            last_err = RuntimeError(f"{exc.code} {exc.read().decode('utf-8', errors='replace')[:300]}")
            if exc.code in {429, 500, 502, 503, 504} and attempt < 3:
                time.sleep(2 * attempt)
                continue
            raise last_err from exc
        except Exception as exc:
            last_err = exc
            if attempt < 3:
                time.sleep(2 * attempt)
                continue
            raise
    raise RuntimeError(f"Gemini failed after retries: {last_err}")


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

def main() -> int:
    spec_path = RESULTS / "spec_gpt.json"
    if not spec_path.exists():
        print(f"[ERROR] {spec_path} not found. Run test_01 first.")
        return 1
    spec = json.loads(spec_path.read_text())

    dsl_text = spec_to_natural_dsl(spec)
    (RESULTS / "spec_gpt_natural_dsl.txt").write_text(dsl_text, encoding="utf-8")
    print(f"[INFO] DSL: {len(dsl_text)}자 → spec_gpt_natural_dsl.txt")

    prompt = build_variant_d_prompt(dsl_text)
    (RESULTS / "variant_d_prompt.txt").write_text(prompt, encoding="utf-8")
    print(f"[INFO] variant D prompt: {len(prompt)}자 → variant_d_prompt.txt")

    print("\n[INFO] gpt-image-2 + Gemini nano-banana-2 병렬 호출...")
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as ex:
        f_gpt = ex.submit(call_gpt_image_2, prompt)
        f_gem = ex.submit(call_gemini_nano_banana, prompt)
        results: Dict[str, Tuple[str, bytes | str]] = {}
        for name, fut in (("gpt-image-2", f_gpt), ("nano-banana-2", f_gem)):
            try:
                img_bytes = fut.result(timeout=600)
                results[name] = ("ok", img_bytes)
            except Exception as exc:
                results[name] = ("error", f"{type(exc).__name__}: {exc}")

    for name, (status, payload) in results.items():
        slug = name.replace("-", "_").replace(".", "")
        if status == "ok":
            out = RESULTS / f"variant_d_{slug}.png"
            out.write_bytes(payload)
            print(f"[OK] {name} → {out} ({len(payload):,} bytes)")
        else:
            err_path = RESULTS / f"variant_d_{slug}_error.txt"
            err_path.write_text(str(payload), encoding="utf-8")
            print(f"[ERROR] {name}: {payload}")

    return 0


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