"""Spike Test 4 (Variant E) — ASCII grid layout in prompt.

가설: 단순한 2D char grid (예: 'T'=TV, 'S'=sofa)가 SVG/DSL보다 모델이 직관적으로
      이해. 학습 데이터에 게임 맵, 코드 댓글의 ASCII art 풍부.

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

CELL_SIZE = 0.5  # meters per cell


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


# ───────── char codes ─────────

ZONE_CHARS = {
    "living": "L",
    "kitchen": "L",  # combined
    "main_bedroom": "M",
    "daughter_bedroom": "D",
    "bathroom": "B",
    "entryway": "E",
}

FURN_CHARS = {
    "tv": "T",
    "sofa": "S",
    "table": "t",
    "sink": "K",
    "stove": "k",
    "refrigerator": "f",
    "bed": "b",
    "mirror": "m",
    "toilet": "o",
    "curtain": "c",
    "chair": ",",
    "shelf": "h",
    "wardrobe": "w",
}

DOOR_CHARS = {
    "front_entry": "F",  # 빨간 문, 외부로
    "interior": "I",     # 초록 문, 방 사이
    "window": "w",
}

EXTERIOR_CHARS = {
    "rooftop_concrete": "R",
    "balcony": "?",
    "street": "?",
}

EMPTY_FLOOR = "."  # zone 안 빈 바닥
OUTSIDE = " "      # apartment 밖, exterior zone도 아닌 곳


# ───────── point in polygon ─────────

def _point_in_polygon(point, polygon) -> bool:
    x, y = point
    inside = False
    n = len(polygon)
    j = n - 1
    for i in range(n):
        xi, yi = polygon[i]
        xj, yj = polygon[j]
        if ((yi > y) != (yj > y)) and (x < (xj - xi) * (y - yi) / (yj - yi + 1e-12) + xi):
            inside = not inside
        j = i
    return inside


# ───────── spec.json → ASCII grid ─────────

def spec_to_grid(spec: Dict[str, Any], cell_size: float = CELL_SIZE) -> Tuple[str, Dict[str, str]]:
    """spec.json → ASCII grid + legend.

    각 cell이 zone char 또는 furniture char.
    가구가 zone 위에 덮어씀 (가구가 우선).
    door는 가구보다 우선.
    grid는 north-up (top row = +y 큰 쪽).
    """
    canvas = spec.get("canvas", {"width": 10, "height": 10})
    cw = float(canvas["width"])
    ch = float(canvas["height"])

    # extend bbox to include exterior
    xmin, ymin, xmax, ymax = 0.0, 0.0, cw, ch
    for ez in spec.get("exterior_adjacency_zones", []):
        for x, y in ez.get("polygon", []):
            xmin = min(xmin, x)
            ymin = min(ymin, y)
            xmax = max(xmax, x)
            ymax = max(ymax, y)

    cols = max(1, int((xmax - xmin) / cell_size) + 1)
    rows = max(1, int((ymax - ymin) / cell_size) + 1)

    # init: 모두 OUTSIDE
    grid = [[OUTSIDE for _ in range(cols)] for _ in range(rows)]

    def grid_xy(r: int, c: int) -> Tuple[float, float]:
        return (
            xmin + (c + 0.5) * cell_size,
            ymin + (r + 0.5) * cell_size,
        )

    # 1) exterior zones (R = rooftop)
    for ez in spec.get("exterior_adjacency_zones", []):
        char = EXTERIOR_CHARS.get(ez.get("type", ""), "X")
        poly = ez.get("polygon", [])
        if len(poly) < 3:
            continue
        for r in range(rows):
            for c in range(cols):
                if _point_in_polygon(grid_xy(r, c), poly):
                    grid[r][c] = char

    # 2) rooms (zone char)
    for room in spec.get("rooms", []):
        rid = room["id"]
        char = ZONE_CHARS.get(rid, rid[0].upper() if rid else "?")
        poly = room.get("polygon", [])
        if len(poly) < 3:
            continue
        for r in range(rows):
            for c in range(cols):
                if _point_in_polygon(grid_xy(r, c), poly):
                    grid[r][c] = char

    # 3) furniture (zone 위에 덮어씀)
    for f in spec.get("furniture", []):
        fx, fy = f.get("position", [0, 0])
        sw, sh = f.get("size", [cell_size, cell_size])
        char = FURN_CHARS.get(f.get("type", ""), f.get("type", "?")[0].lower())
        # 가구 bounding box 안의 cell들 모두 furniture char
        c_start = max(0, int((fx - sw / 2 - xmin) / cell_size))
        c_end = min(cols - 1, int((fx + sw / 2 - xmin) / cell_size))
        r_start = max(0, int((fy - sh / 2 - ymin) / cell_size))
        r_end = min(rows - 1, int((fy + sh / 2 - ymin) / cell_size))
        for r in range(r_start, r_end + 1):
            for c in range(c_start, c_end + 1):
                grid[r][c] = char

    # 4) doors (가장 우선)
    for d in spec.get("doors", []):
        x, y = d.get("position", [0, 0])
        char = DOOR_CHARS.get(d.get("kind", "interior"), "?")
        c = int((x - xmin) / cell_size)
        r = int((y - ymin) / cell_size)
        if 0 <= r < rows and 0 <= c < cols:
            grid[r][c] = char

    # 출력 (north up, row 0 = bottom 이므로 reverse)
    lines = []
    for r in reversed(range(rows)):
        lines.append("".join(grid[r]))

    # legend 빌드
    used_chars = set()
    for line in lines:
        used_chars.update(line)

    legend_pool: Dict[str, str] = {}
    for rid, ch_ in ZONE_CHARS.items():
        if ch_ in used_chars:
            legend_pool[ch_] = f"{rid.replace('_', ' ')} (zone interior)"
    for ftype, ch_ in FURN_CHARS.items():
        if ch_ in used_chars:
            legend_pool[ch_] = f"{ftype}"
    for dkind, ch_ in DOOR_CHARS.items():
        if ch_ in used_chars:
            legend_pool[ch_] = f"{dkind.replace('_', ' ')}"
    for etype, ch_ in EXTERIOR_CHARS.items():
        if ch_ in used_chars:
            legend_pool[ch_] = f"{etype.replace('_', ' ')} (outside the apartment)"
    if EMPTY_FLOOR in used_chars:
        legend_pool[EMPTY_FLOOR] = "empty floor inside a zone"
    if OUTSIDE in used_chars:
        legend_pool[OUTSIDE] = "outside the building / unspecified"

    return "\n".join(lines), legend_pool


# ───────── Build variant E prompt ─────────

def build_variant_e_prompt(grid_text: str, legend: Dict[str, str]) -> str:
    legend_lines = []
    for ch_ in sorted(legend.keys()):
        rendered = "(space)" if ch_ == " " else f"'{ch_}'"
        legend_lines.append(f"  {rendered} = {legend[ch_]}")
    legend_block = "\n".join(legend_lines)

    rows_count = len(grid_text.splitlines())
    cols_count = len(grid_text.splitlines()[0]) if rows_count else 0

    return (
        "[ARCHITECTURAL FLOOR PLAN — ASCII grid, top-down view, used as authoritative spatial layout]\n"
        f"Each cell = {CELL_SIZE}m × {CELL_SIZE}m. North is up (top row), East is right (right column). "
        f"Grid is {cols_count} cols × {rows_count} rows.\n\n"
        "Legend (each character represents what occupies that cell):\n"
        f"{legend_block}\n\n"
        "Layout (read top-down, north up):\n"
        "```\n"
        f"{grid_text}\n"
        "```\n\n"
        "Use this grid as the authoritative spatial layout. The relative "
        "positions of TV (T), sofa (S), kitchen sink (K), table (t), windows "
        "and doors must match this grid exactly.\n\n"
        "[RENDER INSTRUCTIONS]\n"
        + CHAIN_BG_PROMPT
        + "\n\nRender as a single photorealistic eye-level cinematic frame "
        "from the camera position described in the prompt above (NOT a "
        "top-down view, NOT an ASCII art). Use the grid only to determine "
        "where each piece of furniture and each door is — then render it "
        "as a normal interior photograph."
    )


# ───────── 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.")
        return 1
    spec = json.loads(spec_path.read_text())

    grid_text, legend = spec_to_grid(spec)
    (RESULTS / "spec_gpt_grid.txt").write_text(grid_text, encoding="utf-8")

    legend_dump = "\n".join(f"{k}: {v}" for k, v in sorted(legend.items()))
    (RESULTS / "spec_gpt_grid_legend.txt").write_text(legend_dump, encoding="utf-8")

    print(f"[INFO] grid: {len(grid_text.splitlines())} rows × "
          f"{len(grid_text.splitlines()[0])} cols → spec_gpt_grid.txt")
    print(f"\n--- GRID PREVIEW ---")
    print(grid_text)
    print(f"\n--- LEGEND ---")
    print(legend_dump)

    prompt = build_variant_e_prompt(grid_text, legend)
    (RESULTS / "variant_e_prompt.txt").write_text(prompt, encoding="utf-8")
    print(f"\n[INFO] variant E prompt: {len(prompt)}자 → variant_e_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_e_{slug}.png"
            out.write_bytes(payload)
            print(f"[OK] {name} → {out} ({len(payload):,} bytes)")
        else:
            err_path = RESULTS / f"variant_e_{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())