"""실험 — 씬 인지 line elevation + photo 4분할.

가설: 도면(top-down) + 그 도면이 속한 모든 씬/샷의 wall-level 정보를 LLM에 함께 주면,
4면(N/E/S/W) 각각에 대한 정확한 description + 그것에 부합하는 line/photo t2i를 만들 수 있다.

흐름:
  Step A. run_dir + base_plan_id에서 도면 + base_photo + 관련 shots 전수 수집
  Step B. GPT vision: 도면 PNG + 위 텍스트 → JSON {wall_descriptions, line_elevation_t2i, photo_quad_t2i}
  Step C. line_elevation_t2i + 도면 → line PNG
  Step D. photo_quad_t2i + line + 도면 → photo PNG

사용:
  python scripts/experiment_scene_aware_line_quad.py \
    --plan-png <path> --run-dir <v4 run> --base-plan-id <id> --out-dir <dir>
"""
from __future__ import annotations

import argparse
import base64
import glob
import json
import logging
import os
import sys
from pathlib import Path

BACKEND = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(BACKEND))

from dotenv import load_dotenv  # noqa: E402
load_dotenv(BACKEND / ".env")

from openai import OpenAI  # noqa: E402

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("scene_line_quad")


def collect_scene_context(run_dir: Path, base_plan_id: str) -> dict:
    """Gather floor plan spec + base_photo spec + all shots for this base_plan_id."""
    ctx: dict = {
        "base_plan_id": base_plan_id,
        "base_plan": None,
        "base_photo": None,
        "shots": [],
        "spatial": None,
    }

    sp2 = run_dir / "step2_plan_specs.json"
    if sp2.exists():
        d = json.loads(sp2.read_text(encoding="utf-8"))
        for p in d.get("base_plans", []):
            if p.get("id") == base_plan_id:
                ctx["base_plan"] = p
                break

    sp7 = run_dir / "step7_photo_specs.json"
    if sp7.exists():
        d = json.loads(sp7.read_text(encoding="utf-8"))
        for ph in d.get("base_photos", []):
            if ph.get("source_plan_id") == base_plan_id:
                ctx["base_photo"] = ph
                break

    for f in sorted(glob.glob(str(run_dir / "step3_shot_*.json"))):
        s = json.loads(Path(f).read_text(encoding="utf-8"))
        if s.get("base_plan_id") == base_plan_id:
            ctx["shots"].append(s)

    sp1 = run_dir / "step1_spatial.json"
    if sp1.exists():
        ctx["spatial"] = json.loads(sp1.read_text(encoding="utf-8"))

    return ctx


SYSTEM_PROMPT = """You are an architectural visualization specialist.

You will receive:
1. A top-down architectural FLOOR PLAN of an interior room (image attached).
2. A wall-by-wall description of the room (existing N/E/S/W info).
3. SHOTS — narrative scenes set in this room, with camera positions, door states, lighting, and fixed props.

Your task: produce wall descriptions and TWO image prompts.

OUTPUT A — 4-quadrant LINE ELEVATION drawing
  - 2x2 grid, thin black grid divider.
  - Quadrant assignment: TL=NORTH wall, TR=EAST wall, BL=WEST wall, BR=SOUTH wall.
  - Each quadrant: pure black line on white, orthographic FRONTAL elevation, no shading, no perspective, no color.
  - Show as outlined contours: wall outline, floor line, ceiling line, door openings (arc if swing), window frames with mullions, fixed built-ins, and front-facing silhouettes of any furniture sitting against that wall.
  - Quadrant corner labels (N/E/W/S) in ENGLISH ONLY, small clean sans-serif.

OUTPUT B — 4-quadrant PHOTOREALISTIC composite
  - Same 2x2 layout (TL=N, TR=E, BL=W, BR=S), thin black grid divider.
  - Each quadrant: photorealistic eye-level interior view FACING that wall, reproducing the door/window/furniture silhouettes from the corresponding line-elevation quadrant.
  - Camera at room center, eye-level ~1.6m, slight wide angle ~28mm.
  - Same room, same lighting, just facing 4 different walls — visual identity (wall finish, floor finish, ceiling, lighting, palette) consistent across all 4 quadrants.
  - 35mm cinematic still aesthetic, soft natural daylight + dim domestic practicals, realistic shadows, lived-in details (dust, scuff marks).
  - No people, empty room. No overlaid text other than the corner labels.

STRICT RULES — apply to BOTH prompts:
- ENGLISH ONLY in all rendered labels. No Korean, no Chinese characters anywhere in the image.
- Use COMMON NOUNS only. NEVER use proper names of characters, places, scenes, or productions. Say "a small bedroom door" not "<character name>'s bedroom"; "an iron entry door" not the production-specific term.
- Reflect scene context only via NEUTRAL atmospheric cues (lighting tone, wallpaper wear, lived-in dust). Do NOT include narrative events — no people, no blood, no broken glass, no action.
- Keep door/window/furniture positions consistent between line elevation (A) and photo (B) so they map 1:1.
- Each of the 4 quadrants must show DIFFERENT content (different doors/furniture/fixtures). No two quadrants may look alike.
- Quadrant corner labels MUST be the single letters N / E / W / S in the top-left of each quadrant.

Return STRICT JSON, no markdown fence:
{
  "wall_descriptions": {
    "N": "<2-3 sentences in English describing what is on the NORTH wall, viewed from the room center>",
    "E": "<...>",
    "S": "<...>",
    "W": "<...>"
  },
  "line_elevation_t2i": "<full single-string prompt for OUTPUT A — include the 2x2 layout assignment, per-wall content, label rules, and pure-line-art styling>",
  "photo_quad_t2i": "<full single-string prompt for OUTPUT B — reference the line elevation as silhouette guide, include layout, per-wall photo content, lighting, and styling>"
}
"""


def build_user_prompt(ctx: dict) -> str:
    parts: list[str] = []
    parts.append("== BASE PLAN SPEC ==")
    bp = ctx.get("base_plan") or {}
    parts.append(f"id: {bp.get('id')}")
    parts.append(f"label: {bp.get('label')}")
    parts.append(f"visual_domain: {bp.get('visual_domain')}")
    if bp.get("t2i_prompt"):
        parts.append(f"\nplan_drawing_t2i_prompt:\n{bp['t2i_prompt']}")
    if bp.get("legend"):
        parts.append(f"\nlegend: {json.dumps(bp['legend'], ensure_ascii=False)}")
    if bp.get("elements_meta"):
        parts.append(f"\nelements_meta: {json.dumps(bp['elements_meta'], ensure_ascii=False)}")

    parts.append("\n== BASE PHOTO SPEC (already wall-by-wall annotated) ==")
    bph = ctx.get("base_photo") or {}
    if bph.get("camera_note"):
        parts.append(f"camera_note: {bph['camera_note']}")
    if bph.get("lighting"):
        parts.append(f"lighting: {bph['lighting']}")
    if bph.get("t2i_prompt"):
        parts.append(f"\nbase_photo_t2i (contains NORTH/SOUTH/EAST/WEST wall sentences):\n{bph['t2i_prompt']}")

    parts.append("\n== RELATED SHOTS (each shot in this room) ==")
    for s in ctx.get("shots", []):
        si = s.get("scene_index"); sx = s.get("shot_index")
        parts.append(f"\n--- Shot S{si:02d}_Shot{sx} ---")
        cam = s.get("camera") or {}
        parts.append(f"camera position: {cam.get('position', '(none)')}")
        parts.append(f"camera heading: {cam.get('heading', '(none)')}")
        parts.append(f"camera height/fov/lens: {cam.get('height')} / {cam.get('fov')} / {cam.get('lens_note')}")
        if s.get("characters"):
            parts.append(f"characters: {json.dumps(s['characters'], ensure_ascii=False)}")
        for a in s.get("additions", []):
            parts.append(f"  + [{a.get('type')}] @ {a.get('position', '?')}: {a.get('note', '')}")

    parts.append("\n== INTERIOR CANON ==")
    sp = ctx.get("spatial") or {}
    interior = (sp.get("environment_canon") or {}).get("interior") or {}
    if interior:
        parts.append(json.dumps(interior, ensure_ascii=False, indent=2))

    parts.append("\n== TASK ==")
    parts.append("Use the floor plan IMAGE (attached) and ALL information above to produce:")
    parts.append("1) wall_descriptions: one paragraph per wall (N/E/S/W) describing exactly what is visible on that wall when viewed frontally from the room center.")
    parts.append("2) line_elevation_t2i: a complete prompt for the 4-quadrant LINE ELEVATION drawing (TL=N, TR=E, BL=W, BR=S).")
    parts.append("3) photo_quad_t2i: a complete prompt for the 4-quadrant PHOTOREALISTIC composite that uses the line elevation as silhouette guide.")
    parts.append("\nReturn JSON only. No markdown fence, no commentary outside JSON.")
    return "\n".join(parts)


def call_text_llm(client: OpenAI, model: str, plan_png: Path,
                  system_prompt: str, user_prompt: str) -> dict:
    with open(plan_png, "rb") as f:
        b64 = base64.b64encode(f.read()).decode("ascii")

    resp = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": [
                {"type": "text", "text": user_prompt},
                {"type": "image_url",
                 "image_url": {"url": f"data:image/png;base64,{b64}"}},
            ]},
        ],
        response_format={"type": "json_object"},
    )
    txt = resp.choices[0].message.content or ""
    return json.loads(txt)


def edit_call(client: OpenAI, model: str, refs: list[Path], prompt: str,
              size: str, quality: str, out_path: Path) -> Path:
    logger.info("[image edit] %s + %d ref(s) → %s (%s, %s)",
                model, len(refs), out_path.name, size, quality)
    files = [open(p, "rb") for p in refs]
    try:
        resp = client.images.edit(
            model=model, image=files, prompt=prompt,
            size=size, quality=quality, n=1,
        )
    finally:
        for f in files:
            f.close()
    b64 = resp.data[0].b64_json
    if not b64:
        raise RuntimeError(f"empty b64 for {out_path.name}")
    out_path.write_bytes(base64.b64decode(b64))
    logger.info("  saved %d KB", out_path.stat().st_size // 1024)
    return out_path


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--plan-png", required=True)
    p.add_argument("--run-dir", required=True)
    p.add_argument("--base-plan-id", required=True)
    p.add_argument("--out-dir", required=True)
    p.add_argument("--text-model", default=os.getenv("OPENAI_MODEL", "gpt-5.5"))
    p.add_argument("--image-model", default="gpt-image-2")
    p.add_argument("--line-size", default="1024x1024")
    p.add_argument("--photo-size", default="1024x1024")
    p.add_argument("--quality", default="high")
    p.add_argument("--skip-text", action="store_true",
                   help="기존 prompts.json 재사용 (이미지만 재생성)")
    p.add_argument("--skip-line", action="store_true",
                   help="line PNG도 기존 사용 (photo만 재생성)")
    p.add_argument("--text-only", action="store_true",
                   help="LLM 텍스트 단계만 실행 (이미지 생성 없이 prompt 검토용)")
    args = p.parse_args()

    plan_png = Path(args.plan_png).resolve()
    run_dir = Path(args.run_dir).resolve()
    out_dir = Path(args.out_dir).resolve()
    out_dir.mkdir(parents=True, exist_ok=True)

    if not os.getenv("OPENAI_API_KEY"):
        logger.error("OPENAI_API_KEY not set"); return 1
    client = OpenAI()

    prompts_p = out_dir / "prompts.json"

    # Step A + B: collect + LLM
    if args.skip_text and prompts_p.exists():
        prompts = json.loads(prompts_p.read_text(encoding="utf-8"))
        logger.info("[skip text] reusing existing prompts.json")
    else:
        ctx = collect_scene_context(run_dir, args.base_plan_id)
        if not ctx["base_plan"]:
            logger.error("base_plan_id 매칭 실패: %s", args.base_plan_id); return 1
        logger.info("collected: base_plan=ok, base_photo=%s, %d shots",
                    bool(ctx["base_photo"]), len(ctx["shots"]))
        user_prompt = build_user_prompt(ctx)
        logger.info("user prompt: %d chars", len(user_prompt))

        prompts = call_text_llm(client, args.text_model, plan_png,
                                SYSTEM_PROMPT, user_prompt)
        prompts["_meta"] = {
            "base_plan_id": args.base_plan_id,
            "shot_count": len(ctx["shots"]),
            "text_model": args.text_model,
            "user_prompt_chars": len(user_prompt),
        }
        prompts_p.write_text(
            json.dumps(prompts, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        logger.info("saved prompts.json")

    # Print for review
    print("\n" + "=" * 60)
    print("=== WALL DESCRIPTIONS ===")
    print("=" * 60)
    for d in ("N", "E", "S", "W"):
        wd = (prompts.get("wall_descriptions") or {}).get(d, "(missing)")
        print(f"\n[{d}]\n{wd}")
    print("\n" + "=" * 60)
    print("=== LINE ELEVATION T2I ===")
    print("=" * 60)
    print(prompts.get("line_elevation_t2i") or "(missing)")
    print("\n" + "=" * 60)
    print("=== PHOTO QUAD T2I ===")
    print("=" * 60)
    print(prompts.get("photo_quad_t2i") or "(missing)")
    print("=" * 60 + "\n")

    if args.text_only:
        logger.info("text-only mode → skip image generation")
        return 0

    # Step C: line elevation
    line_p = out_dir / "line_elevation_quad.png"
    line_t2i = prompts.get("line_elevation_t2i", "")
    if not line_t2i:
        logger.error("line_elevation_t2i missing"); return 1
    if args.skip_line and line_p.exists():
        logger.info("[skip line] reusing %s", line_p.name)
    else:
        edit_call(client, args.image_model, [plan_png], line_t2i,
                  args.line_size, args.quality, line_p)

    # Step D: photo from line + plan
    photo_p = out_dir / "photo_quad_from_line.png"
    photo_t2i = prompts.get("photo_quad_t2i", "")
    if not photo_t2i:
        logger.error("photo_quad_t2i missing"); return 1
    edit_call(client, args.image_model, [line_p, plan_png], photo_t2i,
              args.photo_size, args.quality, photo_p)

    logger.info("=== DONE === %s", out_dir)
    return 0


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