"""실험 Phase 2 — chain_structure.json을 따라 노드별 t2i prompt 생성 + image chain 렌더링.

흐름 (execution_order대로 노드 1개씩 처리):
  Step 1. LLM call (Gemini Pro, 텍스트 only):
            노드 description + 부모 description + group canon + shared_visual_anchors
            → 그 노드용 t2i prompt 1개
  Step 2. Banana(image gen):
            - 루트 anchor: 도면 PNG ref + 노드 t2i prompt → image
            - 자식 노드: 부모의 렌더된 image ref + 노드 t2i prompt → image

각 노드마다 LLM 호출 1번 + image 호출 1번. 한 번에 다 묻지 않음.
"""
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("chain_render")

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

PROMPT_GEN_SYSTEM = """You write a single photorealistic image-generation prompt for ONE node in an image chain.

Pipeline:
- Each node is rendered as one photo.
- Root anchors get the architectural floor plan as the reference image.
- Non-root nodes get the previously-rendered PARENT photo as the reference image.
- The chain's consistency depends on the prompt's atmosphere/material/lighting wording reusing the parent's identity, and on explicit mention of the shared visual anchors that overlap with the parent.

Hard constraints in your output prompt:
- 35mm cinematic still aesthetic.
- Eye-level approx 1.6m, slight wide angle ~28mm (unless the original shot's camera info contradicts).
- ENGLISH ONLY anywhere in the rendered image.
- COMMON NOUNS only — never use proper names of characters, places, or productions.
- NO people, NO blood, NO broken glass, NO action. Reflect atmosphere only via worn surfaces, dust, dim light, etc.
- For non-root nodes: open the prompt with an atmospheric anchor like "Same room/space as the previous reference image — match its wall finish, floor, ceiling, lighting tone, and color palette". Then explicitly mention each of the shared_visual_anchors_with_parent (e.g. "the small bedroom doorway frame", "the textured fabric curtain") so the image model reuses them.
- For root anchors: thoroughly describe wall finish, floor, ceiling, lighting fixtures, palette — no parent photo to inherit from, so the prompt must be self-sufficient.
- Be concrete: name doors, windows, fixtures, furniture in frame. Avoid abstract narrative descriptions.

Return strict JSON: {"t2i_prompt": "<single English string>"}.
No markdown fence, no extra commentary.
"""

PROMPT_GEN_SCHEMA = {
    "type": "object",
    "properties": {"t2i_prompt": {"type": "string"}},
    "required": ["t2i_prompt"],
}


def collect_context(run_dir: Path, base_plan_ids: list[str]) -> dict:
    bp_set = set(base_plan_ids)
    ctx: dict = {"base_plans": [], "base_photos": [], "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") in bp_set:
                ctx["base_plans"].append(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_node_meta(plan: dict, ctx: dict) -> dict:
    """node_id → original source dict for prompt builder."""
    out: dict = {}
    for n in plan.get("nodes", []):
        nid = n["id"]
        kind = n.get("kind")
        src = n.get("source") or ""
        if kind == "shot":
            for s in ctx["shots"]:
                sid = f"S{s['scene_index']:02d}_Shot{s['shot_index']}"
                if sid == src or sid == nid:
                    out[nid] = {"kind": "shot", "shot": s}; break
        elif kind == "base_photo":
            for ph in ctx["base_photos"]:
                if ph.get("id") == src or ph.get("id") == nid:
                    out[nid] = {"kind": "base_photo", "photo": ph}; break
    return out


def find_anchor_plan_png(node: dict, plan: dict, ctx: dict, run_dir: Path) -> Path | None:
    bp_id: str | None = None
    if node.get("kind") == "base_photo":
        for ph in ctx["base_photos"]:
            if ph.get("id") in (node.get("source"), node.get("id")):
                bp_id = ph.get("source_plan_id"); break
    elif node.get("kind") == "anchor_synthetic":
        for g in plan.get("groups", []):
            if node["id"] in g.get("node_ids", []):
                gn = g.get("name", "")
                for bp in ctx["base_plans"]:
                    if bp["id"] in gn or gn in bp["id"]:
                        bp_id = bp["id"]; break
                break
    if bp_id:
        p = run_dir / f"base_plan_{bp_id}.png"
        if p.exists():
            return p
    return None


def build_per_node_user_prompt(node: dict, parent: dict | None, group: dict | None,
                               meta: dict | None) -> str:
    parts: list[str] = []
    parts.append("== NODE TO GENERATE PROMPT FOR ==")
    parts.append(f"id: {node['id']}")
    parts.append(f"kind: {node['kind']}")
    parts.append(f"label: {node.get('label', '')}")
    parts.append(f"description: {node.get('description', '')}")
    sva = node.get("shared_visual_anchors_with_parent") or []
    parts.append(f"shared_visual_anchors_with_parent: {json.dumps(sva, ensure_ascii=False)}")

    if parent:
        parts.append("\n== PARENT NODE ==")
        parts.append(f"id: {parent['id']}")
        parts.append(f"label: {parent.get('label', '')}")
        parts.append(f"description: {parent.get('description', '')}")
        parts.append("\n[Pipeline note: parent's already-rendered photo will be the reference image. The prompt must reuse the parent's atmosphere/material/lighting language and explicitly mention each shared_visual_anchor.]")
    else:
        parts.append("\n[ROOT ANCHOR. Reference image = floor plan only. Prompt must thoroughly describe wall/floor/ceiling/lighting/palette since there is no prior rendered photo.]")

    if group:
        parts.append("\n== GROUP CANON ==")
        parts.append(f"name: {group.get('name', '')}")
        parts.append(f"shared_canon: {group.get('shared_canon', '')}")

    if meta:
        kind = meta.get("kind")
        if kind == "shot":
            shot = meta["shot"]
            parts.append("\n== ORIGINAL SHOT INFO ==")
            cam = shot.get("camera") or {}
            parts.append(f"camera position: {cam.get('position', '')}")
            parts.append(f"camera heading: {cam.get('heading', '')}")
            parts.append(f"camera height/fov/lens: {cam.get('height')} / {cam.get('fov')} / {cam.get('lens_note')}")
            for a in shot.get("additions", []):
                parts.append(f"  + [{a.get('type')}] @ {a.get('position', '')}: {a.get('note', '')}")
        elif kind == "base_photo":
            ph = meta["photo"]
            parts.append("\n== EXISTING BASE PHOTO SPEC (reuse / refine) ==")
            parts.append(f"camera_note: {ph.get('camera_note', '')}")
            parts.append(f"lighting: {ph.get('lighting', '')}")
            if ph.get("t2i_prompt"):
                parts.append(f"existing_t2i_prompt:\n{ph['t2i_prompt']}")

    parts.append("\n== TASK ==")
    parts.append("Write ONE photorealistic t2i prompt for this node only. Return JSON: {\"t2i_prompt\": \"...\"}")
    return "\n".join(parts)


def call_gemini_text_json(model: str, system: str, user: str, schema: dict) -> dict:
    body = {
        "systemInstruction": {"parts": [{"text": system}]},
        "contents": [{"role": "user", "parts": [{"text": user}]}],
        "generationConfig": {
            "temperature": 0.2,
            "responseMimeType": "application/json",
            "responseJsonSchema": 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=180) 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 {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 empty response: {json.dumps(payload)[:600]}")
    return json.loads(text)


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--run-dir", required=True)
    p.add_argument("--base-plan-ids", required=True)
    p.add_argument("--planning-dir", required=True,
                   help="Phase 1 결과 디렉토리 (chain_structure.json 위치)")
    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("--photo-aspect", default="16:9")
    p.add_argument("--prompts-only", action="store_true",
                   help="t2i prompt 생성만, 이미지 단계 skip")
    p.add_argument("--skip-prompts", action="store_true",
                   help="기존 node_prompts.json 재사용")
    args = p.parse_args()

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

    plan_path = planning_dir / "chain_structure.json"
    if not plan_path.exists():
        logger.error("chain_structure.json 없음: %s", plan_path); return 1
    plan = json.loads(plan_path.read_text(encoding="utf-8"))

    base_plan_ids = [s.strip() for s in args.base_plan_ids.split(",") if s.strip()]
    ctx = collect_context(run_dir, base_plan_ids)
    logger.info("collected: %d base_plans, %d base_photos, %d shots",
                len(ctx["base_plans"]), len(ctx["base_photos"]), len(ctx["shots"]))

    if not os.getenv("GEMINI_API_KEY"):
        logger.error("GEMINI_API_KEY not set"); return 1

    nodes_by_id = {n["id"]: n for n in plan.get("nodes", [])}
    groups_by_node: dict = {}
    for g in plan.get("groups", []):
        for nid in g.get("node_ids", []):
            groups_by_node[nid] = g
    meta_by_node = build_node_meta(plan, ctx)
    execution_order = plan.get("execution_order", [])

    prompts_p = out_dir / "node_prompts.json"

    # === Step 1: 노드별 t2i prompt 생성 ===
    if args.skip_prompts and prompts_p.exists():
        rendered_prompts: dict[str, str] = json.loads(prompts_p.read_text(encoding="utf-8"))
        logger.info("[skip prompts] reusing node_prompts.json (%d entries)", len(rendered_prompts))
    else:
        rendered_prompts = {}
        for i, nid in enumerate(execution_order, 1):
            n = nodes_by_id.get(nid)
            if not n:
                logger.error("missing node id in execution_order: %s", nid); continue
            parent = nodes_by_id.get(n.get("parent_id")) if n.get("parent_id") else None
            group = groups_by_node.get(nid)
            meta = meta_by_node.get(nid)
            user = build_per_node_user_prompt(n, parent, group, meta)
            logger.info("[prompt %d/%d] node=%s (parent=%s) user_prompt=%d chars",
                        i, len(execution_order), nid,
                        (parent or {}).get("id") or "(root)", len(user))
            try:
                resp = call_gemini_text_json(args.text_model,
                                              PROMPT_GEN_SYSTEM, user, PROMPT_GEN_SCHEMA)
                rendered_prompts[nid] = resp.get("t2i_prompt", "")
            except Exception as e:
                logger.error("prompt gen failed for %s: %s", nid, e); continue
        prompts_p.write_text(json.dumps(rendered_prompts, ensure_ascii=False, indent=2),
                             encoding="utf-8")
        logger.info("saved node_prompts.json")

    print("\n" + "=" * 70)
    print("=== PER-NODE T2I PROMPTS ===")
    print("=" * 70)
    for nid in execution_order:
        n = nodes_by_id.get(nid, {})
        print(f"\n--- {nid}  [{n.get('kind')}, parent={n.get('parent_id') or '(root)'}] ---")
        print(rendered_prompts.get(nid, "(missing)"))
    print()

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

    # === Step 2: 노드별 image 생성 (execution_order대로) ===
    image_client = GeminiImageClient(model=args.image_model)
    rendered_images: dict[str, Path] = {}

    for i, nid in enumerate(execution_order, 1):
        n = nodes_by_id.get(nid)
        if not n:
            continue
        prompt_text = rendered_prompts.get(nid)
        if not prompt_text:
            logger.error("missing prompt for %s — skip image", nid); continue

        labeled_refs: list = []
        parent_id = n.get("parent_id")
        if parent_id and parent_id not in ("null", "None", ""):
            parent_path = rendered_images.get(parent_id)
            if not parent_path:
                logger.error("parent image missing for %s (parent=%s) — skip",
                             nid, parent_id); continue
            labeled_refs.append((
                f"Parent rendered photo (consistency anchor — match its wall finish, floor, ceiling, lighting, palette):",
                parent_path.read_bytes(),
            ))
        else:
            plan_png = find_anchor_plan_png(n, plan, ctx, run_dir)
            if plan_png:
                labeled_refs.append((
                    "Floor plan reference (top-down spatial guide):",
                    plan_png.read_bytes(),
                ))
            else:
                logger.warning("no plan png found for root anchor %s", nid)

        out_p = out_dir / f"{nid}.png"
        logger.info("[image %d/%d] node=%s with %d ref(s)",
                    i, len(execution_order), nid, len(labeled_refs))
        try:
            img_bytes, ms = image_client.generate_image(
                prompt=prompt_text,
                labeled_references=labeled_refs,
                aspect_ratio=args.photo_aspect,
            )
            out_p.write_bytes(img_bytes)
            rendered_images[nid] = out_p
            logger.info("  saved %s (%d KB, %d ms)", out_p.name,
                        len(img_bytes) // 1024, ms)
        except Exception as e:
            logger.error("image gen failed for %s: %s", nid, e)

    logger.info("=== DONE === %s (rendered %d images)",
                out_dir, len(rendered_images))
    return 0


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