"""실험 Phase 1 — 옥탑방 관련 샷 클러스터링 + image chain 트리/링크드리스트 구조 LLM 설계.

Phase 1만: 구조 설계 / 이미지 생성 없음.

흐름:
  Step A. run_dir에서 지정한 base_plan_ids에 속하는 모든 base_plan + base_photo + shots 수집
  Step B. Gemini Pro vision: 도면 PNG들 + 모든 샷/사진 정보 → JSON
            {
              rationale_summary,
              groups: [{name, shared_canon, node_ids}],
              nodes: [{id, kind, source, label, description, parent_id, depth,
                       rationale, shared_visual_anchors_with_parent}],
              execution_order: [...]   # topological — 부모 먼저
            }
  Step C. JSON 저장 + tree.txt(ASCII) + diagram.mmd(mermaid) 출력

Phase 2 (다른 스크립트)에서 이 JSON을 읽어 실제 t2i prompt 생성 + image chain 렌더링.
"""
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_key_pool import get_next_key  # noqa: E402

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

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


SYSTEM_PROMPT = """You are a visual continuity planner for a film background image pipeline.

You will receive:
1. ONE OR MORE top-down architectural FLOOR PLANS (images attached, each labeled with its base_plan_id).
2. The full SHOT inventory for the same set of base_plans — for each shot: source_plan_id, camera (position/heading/height/fov/lens), narrative additions (door states, lighting notes, fixed props, marker zones, etc.).
3. The base_photo specs (already-defined wide anchor views per base_plan).

Your job is to design an IMAGE-CHAIN structure that maximizes cross-shot visual consistency when each background image is generated using the PARENT NODE's image as the reference.

Pipeline constraints (READ CAREFULLY):
- Each leaf/intermediate node will be rendered as ONE photorealistic image.
- The node's image-gen request takes:
    * the PARENT node's already-rendered image as the reference
    * a t2i prompt that describes this node's view + reuses parent's atmosphere/material/lighting
- Anchor nodes (kind="anchor") have no parent — they are rendered first, text-only.
- Therefore parent-child pairs must share strong visual anchors (a door, a wall, a furniture piece, a lighting fixture, a wallpaper texture) so the chain holds.

Group / cluster heuristics:
- Group all shots that show the SAME spatial subset (e.g., "small bedroom interior", "living-kitchen entry corner", "rooftop terrace overlook"). One group = one connected sub-tree.
- A group's anchor is the most distinctive WIDE view. If a base_photo already exists for that base_plan, reuse it as the anchor (kind="base_photo"). Otherwise propose a synthetic anchor (kind="anchor_synthetic") with a label and description.
- Within a group, prefer SHALLOW trees (depth <= 2) over deep linked-lists. Deep chains accumulate drift.
- A shot whose camera is INSIDE a sub-room (e.g., reveal through a doorway) should be parented to the corresponding doorway-anchor shot, not the wide overview, because the doorway frame is the strongest shared anchor.
- A shot that is a CLOSE-UP of a prop already visible in another shot should be parented to that other shot.
- Different base_plans (interior vs. rooftop terrace exterior) form SEPARATE trees — no cross-plan parenting unless one shot truly contains both (rare).

For EACH node output:
- id: unique short id. Use the original shot id (e.g., "S12_Shot4") for shot nodes; for synthetic anchors, use a snake_case id like "anchor_interior_wide".
- kind: one of "anchor_synthetic" | "base_photo" | "shot"
- source: original shot_id (for kind=shot) or base_photo_id (for kind=base_photo) or null (for kind=anchor_synthetic)
- label: short human-readable label, English
- description: 2-3 sentences in English describing what this node's image should show
- parent_id: id of parent node, or null if this is a root anchor
- depth: 0 for root, 1 for child of root, etc.
- rationale: ONE sentence explaining why this parent was chosen — what visual element will be reused from parent's image
- shared_visual_anchors_with_parent: list of short noun phrases (e.g., ["small bedroom doorway", "wall finish", "ceiling lighting fixture"]) that MUST appear in BOTH parent and child images so the chain stays consistent. Empty list for root anchors.

Also output:
- groups: list of {name (snake_case), shared_canon (one paragraph describing material/lighting/palette canon for this group), node_ids}
- execution_order: topological ordering — every parent strictly before its children. Roots first, then BFS-wise descendants.
- rationale_summary: ONE paragraph explaining the overall structure choices.

STRICT RULES:
- ENGLISH ONLY in all labels, descriptions, rationales.
- COMMON NOUNS only — NEVER use proper names of characters, places, or productions.
- No narrative events in descriptions — focus on visible space/material/light only. No people, no blood.
- Every shot in the input MUST appear as a node. No shot left unplaced.
- Every node (except roots) MUST have parent_id set to an existing node id.
- execution_order MUST contain every node id exactly once, with parents before children.

Return strict JSON. No markdown fence.
"""


RESPONSE_SCHEMA = {
    "type": "object",
    "properties": {
        "rationale_summary": {"type": "string"},
        "groups": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "name": {"type": "string"},
                    "shared_canon": {"type": "string"},
                    "node_ids": {"type": "array", "items": {"type": "string"}},
                },
                "required": ["name", "shared_canon", "node_ids"],
            },
        },
        "nodes": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "id": {"type": "string"},
                    "kind": {"type": "string", "enum": ["anchor_synthetic", "base_photo", "shot"]},
                    "source": {"type": "string"},
                    "label": {"type": "string"},
                    "description": {"type": "string"},
                    "parent_id": {"type": "string"},
                    "depth": {"type": "integer"},
                    "rationale": {"type": "string"},
                    "shared_visual_anchors_with_parent": {
                        "type": "array",
                        "items": {"type": "string"},
                    },
                },
                "required": [
                    "id", "kind", "source", "label", "description",
                    "parent_id", "depth", "rationale",
                    "shared_visual_anchors_with_parent",
                ],
            },
        },
        "execution_order": {"type": "array", "items": {"type": "string"}},
    },
    "required": ["rationale_summary", "groups", "nodes", "execution_order"],
}


def collect_context(run_dir: Path, base_plan_ids: list[str]) -> dict:
    """Gather all base_plans + base_photos + shots whose base_plan_id is in the requested set."""
    bp_set = set(base_plan_ids)
    ctx: dict = {"base_plans": [], "base_photos": [], "shots": [],
                 "spatial": None, "plan_pngs": {}}

    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") in bp_set:
                ctx["base_plans"].append(p)
                # 도면 PNG path
                png_p = run_dir / f"base_plan_{p['id']}.png"
                if png_p.exists():
                    ctx["plan_pngs"][p["id"]] = png_p

    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") in bp_set:
                ctx["base_photos"].append(ph)

    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") in bp_set:
            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 PLANS IN SCOPE ==")
    for p in ctx["base_plans"]:
        parts.append(f"\n--- base_plan id={p['id']} ---")
        parts.append(f"label: {p.get('label')}")
        parts.append(f"visual_domain: {p.get('visual_domain')}")
        parts.append(f"is_anchor: {p.get('is_anchor')}")
        parts.append(f"anchor_cluster_id: {p.get('anchor_cluster_id')}")

    parts.append("\n\n== BASE PHOTOS (existing wide anchors) ==")
    for ph in ctx["base_photos"]:
        parts.append(f"\n--- base_photo id={ph['id']} ---")
        parts.append(f"source_plan_id: {ph.get('source_plan_id')}")
        parts.append(f"label: {ph.get('label')}")
        parts.append(f"visual_domain: {ph.get('visual_domain')}")
        parts.append(f"camera_note: {ph.get('camera_note', '')}")
        parts.append(f"lighting: {ph.get('lighting', '')}")

    parts.append("\n\n== SHOTS ==")
    for s in ctx["shots"]:
        si = s.get("scene_index"); sx = s.get("shot_index")
        sid = f"S{si:02d}_Shot{sx}"
        parts.append(f"\n--- shot id={sid} (base_plan_id={s.get('base_plan_id')}) ---")
        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')}")
        parts.append(f"visual_domain: {s.get('visual_domain')}")
        parts.append(f"characters: {len(s.get('characters', []))}")
        adds = s.get("additions", [])
        if adds:
            parts.append("additions:")
            for a in adds:
                parts.append(f"  + [{a.get('type')}] @ {a.get('position', '?')}: {a.get('note', '')}")

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

    parts.append("\n\n== TASK ==")
    parts.append("Use the floor plan IMAGES (attached, each tagged with its base_plan_id in the text) and ALL shot/photo info above to design the image-chain structure.")
    parts.append("Output JSON with: rationale_summary, groups, nodes (every shot must appear), execution_order (topological).")
    parts.append("Return JSON only. No commentary outside JSON.")
    return "\n".join(parts)


def call_gemini_pro_with_plans(model: str, plan_pngs: dict[str, Path],
                               system_prompt: str, user_prompt: str,
                               response_schema: dict) -> dict:
    user_parts: list = [{"text": user_prompt}]
    for plan_id, png_path in plan_pngs.items():
        user_parts.append({"text": f"\n[Floor plan image follows for base_plan_id={plan_id}]"})
        user_parts.append({
            "inline_data": {
                "mime_type": "image/png",
                "data": base64.b64encode(png_path.read_bytes()).decode("ascii"),
            }
        })

    body = {
        "systemInstruction": {"parts": [{"text": system_prompt}]},
        "contents": [{"role": "user", "parts": user_parts}],
        "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 render_tree_ascii(plan: dict) -> str:
    """ASCII tree from nodes."""
    nodes = {n["id"]: n for n in plan.get("nodes", [])}
    children: dict[str, list[str]] = {}
    roots: list[str] = []
    for n in plan.get("nodes", []):
        pid = n.get("parent_id") or ""
        if not pid or pid not in nodes:
            roots.append(n["id"])
        else:
            children.setdefault(pid, []).append(n["id"])

    lines: list[str] = []

    def walk(nid: str, prefix: str, is_last: bool) -> None:
        n = nodes[nid]
        connector = "└── " if is_last else "├── "
        kind = n.get("kind", "?")
        depth = n.get("depth", "?")
        label = n.get("label", "?")
        lines.append(f"{prefix}{connector}[{kind}] {nid}  d={depth}  — {label}")
        kids = children.get(nid, [])
        new_prefix = prefix + ("    " if is_last else "│   ")
        for i, k in enumerate(kids):
            walk(k, new_prefix, i == len(kids) - 1)

    for i, r in enumerate(roots):
        walk(r, "", i == len(roots) - 1)
    return "\n".join(lines)


def render_mermaid(plan: dict) -> str:
    """Mermaid flowchart syntax for the chain."""
    out = ["flowchart TD"]
    for n in plan.get("nodes", []):
        nid = n["id"]
        kind = n.get("kind", "?")
        label = n.get("label", "?")
        node_label = f"{nid}<br/>{label}"
        if kind == "anchor_synthetic":
            shape = f'(("{node_label}"))'
        elif kind == "base_photo":
            shape = f'[/"{node_label}"/]'
        else:
            shape = f'["{node_label}"]'
        out.append(f"  {nid}{shape}")
    for n in plan.get("nodes", []):
        pid = n.get("parent_id")
        if pid:
            shared = ", ".join(n.get("shared_visual_anchors_with_parent", [])[:3])
            label = f"|{shared}|" if shared else ""
            out.append(f"  {pid} -->{label} {n['id']}")
    return "\n".join(out)


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--run-dir", required=True)
    p.add_argument("--base-plan-ids", required=True,
                   help="콤마 구분, 옥탑방 관련만: 'small_rooftop_room_interior,rooftop_terrace_and_entry'")
    p.add_argument("--out-dir", required=True)
    p.add_argument("--text-model", default=settings.gemini_text_model)
    p.add_argument("--skip-text", action="store_true",
                   help="기존 chain_structure.json 재사용 (시각화만 다시 렌더)")
    args = p.parse_args()

    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

    base_plan_ids = [s.strip() for s in args.base_plan_ids.split(",") if s.strip()]
    plan_p = out_dir / "chain_structure.json"

    if args.skip_text and plan_p.exists():
        plan = json.loads(plan_p.read_text(encoding="utf-8"))
        logger.info("[skip text] reusing chain_structure.json")
    else:
        ctx = collect_context(run_dir, base_plan_ids)
        logger.info("collected: %d base_plans, %d base_photos, %d shots, %d plan_pngs",
                    len(ctx["base_plans"]), len(ctx["base_photos"]),
                    len(ctx["shots"]), len(ctx["plan_pngs"]))
        if not ctx["shots"] and not ctx["base_photos"]:
            logger.error("no shots/photos matched base_plan_ids=%s", base_plan_ids); return 1

        user_prompt = build_user_prompt(ctx)
        logger.info("user prompt: %d chars; calling Gemini Pro %s",
                    len(user_prompt), args.text_model)
        plan = call_gemini_pro_with_plans(
            args.text_model, ctx["plan_pngs"],
            SYSTEM_PROMPT, user_prompt, RESPONSE_SCHEMA,
        )
        plan["_meta"] = {
            "base_plan_ids": base_plan_ids,
            "shot_count": len(ctx["shots"]),
            "base_photo_count": len(ctx["base_photos"]),
            "text_model": args.text_model,
        }
        plan_p.write_text(json.dumps(plan, ensure_ascii=False, indent=2),
                          encoding="utf-8")
        logger.info("saved chain_structure.json")

    # Render tree + mermaid
    tree_txt = render_tree_ascii(plan)
    (out_dir / "tree.txt").write_text(tree_txt + "\n", encoding="utf-8")
    mermaid = render_mermaid(plan)
    (out_dir / "diagram.mmd").write_text(mermaid + "\n", encoding="utf-8")

    print("\n" + "=" * 70)
    print("=== RATIONALE SUMMARY ===")
    print("=" * 70)
    print(plan.get("rationale_summary", "(missing)"))

    print("\n" + "=" * 70)
    print("=== GROUPS ===")
    print("=" * 70)
    for g in plan.get("groups", []):
        print(f"\n• {g.get('name')}: {len(g.get('node_ids', []))} nodes")
        print(f"  shared_canon: {g.get('shared_canon', '')}")
        print(f"  node_ids: {', '.join(g.get('node_ids', []))}")

    print("\n" + "=" * 70)
    print("=== TREE ===")
    print("=" * 70)
    print(tree_txt)

    print("\n" + "=" * 70)
    print("=== EXECUTION ORDER (topological) ===")
    print("=" * 70)
    for i, nid in enumerate(plan.get("execution_order", []), 1):
        n = next((x for x in plan.get("nodes", []) if x.get("id") == nid), {})
        print(f"  {i:2d}. {nid}  [{n.get('kind', '?')}, d={n.get('depth', '?')}]  ← parent={n.get('parent_id') or '(root)'}")

    print("\n" + "=" * 70)
    print("=== NODE DETAILS ===")
    print("=" * 70)
    for n in plan.get("nodes", []):
        print(f"\n[{n['id']}] kind={n.get('kind')} depth={n.get('depth')} parent={n.get('parent_id') or '(root)'}")
        print(f"  label: {n.get('label')}")
        print(f"  description: {n.get('description')}")
        print(f"  rationale: {n.get('rationale')}")
        sva = n.get("shared_visual_anchors_with_parent") or []
        if sva:
            print(f"  shared_visual_anchors: {', '.join(sva)}")

    print(f"\nWrote: {plan_p}")
    print(f"Wrote: {out_dir / 'tree.txt'}")
    print(f"Wrote: {out_dir / 'diagram.mmd'}")
    return 0


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