"""FP 기반 실내 3-stage 스크래치 실험 (2026-07-03, 사용자 구조제안 — 커밋 금지).

구조: fp → (B) fp 위 엔티티+카메라 주입 blocking (I2I)
        → (C) blocking 이 보여주는 카메라의 실내 배경 생성 (I2I, 인물 0)
        → (chain) 다음 배경 생성 시 이전 배경을 ref 로 연결 (일관성)
대조: 같은 Stage C 를 체인 없이 생성 → 육안 비교.

outdoor 3-stage(aerial→blocking→sketch) 와 대칭. 실험 전용 — production 미접촉.
"""
import base64
import json
import os
import sys
from pathlib import Path

ROOT = Path("/Users/manta/Documents/Projects/TheRoad-I1")
OUT = ROOT / "scratchpad" / "fp3_exp"
FP = ROOT / ("projects/b0ad5c18-5140-4022-a2b0-805de0e5a385/episodes/"
             "0d30302a-51c8-4ac9-98e2-6cd9180512bf/images/floor_plan/"
             "fp_l04_rooftop_room_main.png")
CP = ROOT / ("projects/b0ad5c18-5140-4022-a2b0-805de0e5a385/checkpoints/episodes/"
             "0d30302a-51c8-4ac9-98e2-6cd9180512bf")

# .env 로드 (OPENAI_API_KEY)
for line in (ROOT / "backend" / ".env").read_text().splitlines():
    if line.startswith("OPENAI_API_KEY=") and "OPENAI_API_KEY" not in os.environ:
        os.environ["OPENAI_API_KEY"] = line.split("=", 1)[1].strip()

from openai import OpenAI  # noqa: E402
client = OpenAI()
MODEL = "gpt-image-2"

SHOTS = [(14, 4), (14, 8), (5, 3)]   # 침실 wide / 거실 식탁 medium / 주방 싱크 medium

staging = {(s.get("scene_index"), s.get("shot_index")): s
           for s in json.loads((CP / "shot_staging" / "manifest.json").read_text())["data"]["shots"]}


def _img_edit(prompt: str, refs: list, size: str) -> bytes:
    files = [open(p, "rb") for p in refs]
    try:
        r = client.images.edit(model=MODEL, image=files, prompt=prompt, size=size)
    finally:
        for f in files:
            f.close()
    return base64.b64decode(r.data[0].b64_json)


BLOCKING_PROMPT = (
    "The attached image is a TOP-DOWN floor plan of one dwelling: numbered circles mark "
    "fixed rooms and furniture. Keep the plan EXACTLY as it is — same walls, rooms, "
    "furniture and numbered circles, do not redraw or restyle anything.\n"
    "ADD exactly these annotations for ONE camera setup:\n"
    "- ONE green camera icon with a green triangular view-CONE showing where the camera "
    "stands and what it looks at, following this staging description:\n{cam}\n"
    "{figures}"
    "Add NOTHING else — no new rooms, no new furniture, no text other than the existing "
    "numbers and the new capital letters."
)

BG_PROMPT = (
    "The FIRST attached image is a TOP-DOWN floor plan of a dwelling with a green camera "
    "icon and view-cone marking ONE camera position and its viewing direction. Numbered "
    "circles mark rooms and furniture; lettered circles (if any) mark where people stand — "
    "but render NO people.\n"
    "{chain}"
    "Draw the PHOTOREALISTIC interior view that THAT camera sees from its marked position, "
    "at eye level, looking along the view-cone: the rooms, walls, doorways, windows and "
    "furniture that the plan places in front of that camera, at real-world human scale "
    "(doorways adult height, tables waist height). This is an EMPTY-ROOM background plate: "
    "absolutely no people, no text, no markers, no circles, no arrows.\n"
    "A modest, lived-in Korean rooftop flat (oktapbang) interior, natural daylight.\n"
    "Camera staging: {cam}"
)

CHAIN_CLAUSE = (
    "The SECOND attached image is a photo of ANOTHER part of THE SAME dwelling — match "
    "its wall colours, floor material, trim, door style, lighting mood and overall "
    "renovation age EXACTLY, so both photos read as the same home. Do NOT copy its "
    "camera angle or composition.\n"
)


def figures_clause(sh) -> str:
    cons = ((sh.get("frame_spatial_contract") or {}).get("constraints") or [])
    n = sum(1 for c in cons if isinstance(c, dict) and c.get("target_kind") == "character")
    if n <= 0:
        return "- This shot has NO people — add no lettered circles.\n"
    return (f"- {n} red circle(s) with capital letters (A){' (B)' if n > 1 else ''} inside "
            f"the view-cone where the staging places the people.\n")


def main():
    which = sys.argv[1] if len(sys.argv) > 1 else "all"
    prev_bg = None
    for si, shi in SHOTS:
        sh = staging[(si, shi)]
        cam = str(sh.get("camera_direction") or "")
        tag = f"s{si}sh{shi}"

        blk_path = OUT / f"blk_{tag}.png"
        if which in ("all", "blocking") and not blk_path.exists():
            png = _img_edit(
                BLOCKING_PROMPT.format(cam=cam, figures=figures_clause(sh)),
                [FP], "1024x1024")
            blk_path.write_bytes(png)
            print(f"[B] {tag} blocking -> {blk_path.name} ({len(png)}b)", flush=True)

        if which in ("all", "bg"):
            # 체인 버전
            bg_path = OUT / f"bg_chain_{tag}.png"
            if not bg_path.exists():
                refs = [blk_path] + ([prev_bg] if prev_bg else [])
                chain = CHAIN_CLAUSE if prev_bg else ""
                png = _img_edit(BG_PROMPT.format(chain=chain, cam=cam), refs, "1536x1024")
                bg_path.write_bytes(png)
                print(f"[C] {tag} bg(chain={'Y' if prev_bg else 'first'}) -> {bg_path.name}", flush=True)
            prev_bg = bg_path
            # 대조: 체인 없는 버전 (첫 샷은 체인 자체가 없으므로 skip)
            if (si, shi) != SHOTS[0]:
                nc_path = OUT / f"bg_nochain_{tag}.png"
                if not nc_path.exists():
                    png = _img_edit(BG_PROMPT.format(chain="", cam=cam), [blk_path], "1536x1024")
                    nc_path.write_bytes(png)
                    print(f"[C'] {tag} bg(nochain) -> {nc_path.name}", flush=True)
    print("[done]", flush=True)


if __name__ == "__main__":
    main()
