"""실험 — Gemini Pro 분석 + Banana(1.5 Pro)로 도면 재생성(4카메라 표시) + 4 view photo chain.

흐름:
  Step A. run_dir + base_plan_id에서 도면 + base_photo + shots 수집
  Step B. Gemini Pro vision: 도면 + scene → JSON
            {
              redrawn_plan_t2i_prompt,   # 새 도면 — 4 카메라 wedge 표시 포함
              view_0/45/90/135_deg: {description, t2i_prompt},
              atmosphere_canon
            }
  Step C. Banana: 기존 도면 ref + redrawn_plan_t2i_prompt → 새 도면 PNG
  Step D. Banana chain: 새 도면 + 직전 view image + view prompt → 4 view photos
"""
from __future__ import annotations

import argparse
import base64
import glob
import json
import logging
import os
import sys
import time
import urllib.error
import urllib.request
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 app.core.config import settings  # noqa: E402
from app.modules.llm.gemini_image_client import GeminiImageClient  # noqa: E402
from app.modules.llm.gemini_key_pool import get_next_key  # noqa: E402

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

GEMINI_TEXT_URL = (
    "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}"
)

VIEW_KEYS = ["view_0_deg", "view_45_deg", "view_90_deg", "view_135_deg"]
VIEW_DEGREES = [0, 45, 90, 135]


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. Wall-by-wall scene annotation (existing N/E/S/W info).
3. SHOT context — narrative scenes set in this room with camera positions, doors, lighting, props.

Your task — produce ALL of the following:

A) redrawn_plan_t2i_prompt
   A complete English prompt for REDRAWING the floor plan with FOUR CAMERA MARKERS added at the room's geometric center.
   The redrawn plan must:
   - PRESERVE the original layout EXACTLY: every wall, door, window, fixed built-in, and furniture in the same place. Same wall lengths, same room subdivisions, same compass orientation.
   - Same architectural style: black-and-white line drawing on white background, top-down orthographic, no perspective, no color other than the camera markers, no shading.
   - ADD a single shared camera origin marker at the room's geometric center: a small filled circle with label "CAM ORIGIN".
   - ADD 4 separate camera wedge markers radiating from the origin, each in a distinct accent color (e.g. red, orange, magenta, blue) so they read clearly:
       * cam_0   facing 0° (toward the wall labeled NORTH on the plan)
       * cam_45  facing 45° clockwise from cam_0 (toward the north-east corner)
       * cam_90  facing 90° clockwise from cam_0 (toward the wall labeled EAST)
       * cam_135 facing 135° clockwise from cam_0 (toward the south-east corner)
   - Each wedge: a triangular wedge symbol pointing in its direction PLUS a dashed FOV cone (~30° opening) extending from the origin toward the framed wall. Tip of the wedge is the camera origin.
   - English label next to each wedge tip: "0°", "45°", "90°", "135°", in clean sans-serif. Optionally a small 1-2 word note in English (e.g. "→ N wall", "→ NE corner", "→ E wall", "→ SE corner").
   - Preserve the original compass rose (N) and scale bar.
   - Preserve all original ENGLISH labels in the plan (room names, door swing arcs, etc.). NO Korean, NO Chinese, NO other scripts. Common nouns only.
   - The marker overlay should not obscure walls or doors — keep the FOV cones as thin dashed lines and wedges small.

B) Four photorealistic view prompts: view_0_deg, view_45_deg, view_90_deg, view_135_deg
   Each view's prompt MUST:
   - Reference the redrawn floor plan as the spatial guide. Phrase: "the floor plan reference shows four camera wedges at the room center; this image is the photorealistic view from camera <N>° (the wedge labeled <N>°)".
   - Camera at the room's geometric center, eye-level approx 1.6m, slight wide angle ~28mm.
   - Reproduce the door/window/furniture silhouettes that fall within the corresponding camera's FOV cone.
   - Maintain consistent room identity (wall finish, floor, ceiling, lighting, palette) — see atmosphere_canon below.
   - 35mm cinematic still aesthetic, soft natural daylight + dim domestic practicals, realistic shadows.
   - NO people, empty room. NO overlaid text inside the photo.
   - Include description (2-3 English sentences explaining what is in frame) and a single-string t2i_prompt.

C) atmosphere_canon
   ONE paragraph in English summarizing wall finish, floor, ceiling, lighting fixtures, color palette, and ambience that MUST appear identically across all 4 view photos. The image model uses this as a consistency anchor.

STRICT RULES applying to ALL outputs:
- ENGLISH ONLY in any rendered text/label.
- COMMON NOUNS only. NEVER use proper names of characters, places, or productions.
- No narrative events. No people, blood, broken glass, action.
- Each view must show DIFFERENT content (different walls/doors/furniture).
- Mention the corresponding camera marker explicitly in each view's t2i_prompt.

Return strict JSON. No markdown fence.
"""


RESPONSE_SCHEMA = {
    "type": "object",
    "properties": {
        "redrawn_plan_t2i_prompt": {"type": "string"},
        "view_0_deg": {
            "type": "object",
            "properties": {
                "description": {"type": "string"},
                "t2i_prompt": {"type": "string"},
            },
            "required": ["description", "t2i_prompt"],
        },
        "view_45_deg": {
            "type": "object",
            "properties": {
                "description": {"type": "string"},
                "t2i_prompt": {"type": "string"},
            },
            "required": ["description", "t2i_prompt"],
        },
        "view_90_deg": {
            "type": "object",
            "properties": {
                "description": {"type": "string"},
                "t2i_prompt": {"type": "string"},
            },
            "required": ["description", "t2i_prompt"],
        },
        "view_135_deg": {
            "type": "object",
            "properties": {
                "description": {"type": "string"},
                "t2i_prompt": {"type": "string"},
            },
            "required": ["description", "t2i_prompt"],
        },
        "atmosphere_canon": {"type": "string"},
    },
    "required": [
        "redrawn_plan_t2i_prompt",
        "view_0_deg", "view_45_deg", "view_90_deg", "view_135_deg",
        "atmosphere_canon",
    ],
}


def collect_scene_context(run_dir: Path, base_plan_id: str) -> dict:
    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


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_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 (existing wall-by-wall annotation) ==")
    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_prompt:\n{bph['t2i_prompt']}")

    parts.append("\n== RELATED SHOTS ==")
    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"height/fov/lens: {cam.get('height')} / {cam.get('fov')} / {cam.get('lens_note')}")
        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("- redrawn_plan_t2i_prompt (with the 4 camera wedges defined)")
    parts.append("- view_0_deg, view_45_deg, view_90_deg, view_135_deg (each with description + t2i_prompt)")
    parts.append("- atmosphere_canon (consistency anchor)")
    parts.append("Return JSON only.")
    return "\n".join(parts)


def call_gemini_pro_vision(model: str, plan_png_bytes: bytes,
                           system_prompt: str, user_prompt: str,
                           response_schema: dict) -> dict:
    body = {
        "systemInstruction": {"parts": [{"text": system_prompt}]},
        "contents": [{
            "role": "user",
            "parts": [
                {"text": user_prompt},
                {"inline_data": {
                    "mime_type": "image/png",
                    "data": base64.b64encode(plan_png_bytes).decode("ascii"),
                }},
            ],
        }],
        "generationConfig": {
            "temperature": 0.2,
            "responseMimeType": "application/json",
            "responseJsonSchema": response_schema,
            "maxOutputTokens": settings.llm_max_output_tokens,
        },
    }
    last_err: Exception | None = None
    for attempt in range(1, 4):
        api_key = get_next_key()
        url = GEMINI_TEXT_URL.format(model=model, api_key=api_key)
        req = urllib.request.Request(
            url, data=json.dumps(body).encode("utf-8"),
            headers={"Content-Type": "application/json"}, method="POST",
        )
        try:
            with urllib.request.urlopen(req, timeout=240) as resp:
                payload = json.loads(resp.read().decode("utf-8"))
            break
        except urllib.error.HTTPError as e:
            err = e.read().decode("utf-8", errors="replace")
            last_err = RuntimeError(f"Gemini Pro {e.code}: {err[:400]}")
            if e.code in (429, 500, 502, 503, 504) and attempt < 3:
                time.sleep(3 * attempt); continue
            raise last_err
    cand = (payload.get("candidates") or [{}])[0]
    parts = (cand.get("content") or {}).get("parts") or []
    text = "".join(p.get("text", "") for p in parts).strip()
    if not text:
        raise RuntimeError(f"Gemini Pro empty response: {json.dumps(payload)[:600]}")
    return json.loads(text)


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=settings.gemini_text_model)
    p.add_argument("--image-model", default=settings.gemini_image_model)
    p.add_argument("--plan-aspect", default="1:1")
    p.add_argument("--photo-aspect", default="16:9")
    p.add_argument("--skip-text", action="store_true")
    p.add_argument("--skip-plan", action="store_true",
                   help="새 도면 재사용 (이미 만들어진 redrawn_plan.png 사용)")
    p.add_argument("--text-only", action="store_true")
    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("GEMINI_API_KEY"):
        logger.error("GEMINI_API_KEY not set"); return 1

    plan_bytes = plan_png.read_bytes()
    prompts_p = out_dir / "prompts.json"

    # Step A + B: scene context + Gemini Pro
    if args.skip_text and prompts_p.exists():
        prompts = json.loads(prompts_p.read_text(encoding="utf-8"))
        logger.info("[skip text] reusing 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; calling Gemini Pro %s",
                    len(user_prompt), args.text_model)
        prompts = call_gemini_pro_vision(args.text_model, plan_bytes,
                                         SYSTEM_PROMPT, user_prompt, RESPONSE_SCHEMA)
        prompts["_meta"] = {
            "base_plan_id": args.base_plan_id,
            "shot_count": len(ctx["shots"]),
            "text_model": args.text_model,
            "image_model": args.image_model,
            "plan_aspect": args.plan_aspect,
            "photo_aspect": args.photo_aspect,
        }
        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("=== REDRAWN PLAN T2I PROMPT ===")
    print("=" * 60)
    print(prompts.get("redrawn_plan_t2i_prompt", "(missing)"))
    print("\n=== ATMOSPHERE CANON ===")
    print(prompts.get("atmosphere_canon", "(missing)"))
    for k, deg in zip(VIEW_KEYS, VIEW_DEGREES):
        v = prompts.get(k) or {}
        print(f"\n--- {k} ({deg}°) ---")
        print(f"description: {v.get('description', '(missing)')}")
        print(f"t2i_prompt:\n{v.get('t2i_prompt', '(missing)')}")
    print()

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

    image_client = GeminiImageClient(model=args.image_model)

    # Step C: redrawn plan
    new_plan_p = out_dir / "redrawn_plan.png"
    if args.skip_plan and new_plan_p.exists():
        logger.info("[skip plan] reusing %s", new_plan_p.name)
        new_plan_bytes = new_plan_p.read_bytes()
    else:
        plan_t2i = (prompts.get("redrawn_plan_t2i_prompt") or "").strip()
        if not plan_t2i:
            logger.error("redrawn_plan_t2i_prompt missing"); return 1
        logger.info("[redrawn plan] generating with original plan as ref")
        plan_bytes_out, ms = image_client.generate_image(
            prompt=plan_t2i,
            labeled_references=[("Original floor plan reference (preserve layout exactly):", plan_bytes)],
            aspect_ratio=args.plan_aspect,
        )
        new_plan_p.write_bytes(plan_bytes_out)
        new_plan_bytes = plan_bytes_out
        logger.info("  saved %s (%d KB, %d ms)",
                    new_plan_p.name, len(plan_bytes_out) // 1024, ms)

    # Step D: chained view photos using NEW plan as spatial reference
    atmosphere = prompts.get("atmosphere_canon", "")
    prev_image_bytes: bytes | None = None

    for k, deg in zip(VIEW_KEYS, VIEW_DEGREES):
        v = prompts.get(k) or {}
        view_prompt = (v.get("t2i_prompt") or "").strip()
        if not view_prompt:
            logger.error("missing t2i_prompt for %s — skip", k); continue

        full_prompt = view_prompt
        if atmosphere:
            full_prompt += f"\n\nAtmosphere consistency anchor (must match across all views): {atmosphere}"
        if prev_image_bytes is not None:
            full_prompt += (
                "\n\nThe second reference image is the photorealistic view from the previous "
                "rotation step. Match its wall finish, floor finish, ceiling, lighting tone, "
                "color palette, and overall material aging — same room, just rotated. The new "
                "view should show DIFFERENT walls/objects than the previous image."
            )

        labeled_refs = [
            ("Floor plan reference with 4 camera wedges (top-down, this is the spatial guide):", new_plan_bytes),
        ]
        if prev_image_bytes is not None:
            labeled_refs.append(("Previous rotation step photo (consistency reference):", prev_image_bytes))

        out_p = out_dir / f"photo_{k}.png"
        logger.info("[%s @ %d°] generating with %d refs (new_plan%s)",
                    k, deg, len(labeled_refs),
                    " + prev" if prev_image_bytes is not None else "")
        img_bytes, ms = image_client.generate_image(
            prompt=full_prompt,
            labeled_references=labeled_refs,
            aspect_ratio=args.photo_aspect,
        )
        out_p.write_bytes(img_bytes)
        logger.info("  saved %s (%d KB, %d ms)", out_p.name,
                    len(img_bytes) // 1024, ms)
        prev_image_bytes = img_bytes

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


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