"""실험 Phase 2 (GPT 버전) — chain_structure.json을 따라 노드별 t2i prompt 생성 + image chain.

전부 GPT:
  - 노드별 t2i prompt 생성: gpt-5.5 (chat.completions JSON 모드)
  - 이미지 생성: gpt-image-2
      * 루트 anchor: images.generate (텍스트만)  — 도면 ref가 들어가면 더 좋지만 generate는 ref 미지원
        → 루트 anchor도 images.edit + 도면 PNG ref 사용
      * 자식 노드: images.edit + 부모의 이미 렌더된 PNG를 ref로

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


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

INPUT LANGUAGE NOTICE — READ FIRST:
The user message may contain Korean raw scenario text (shot descriptions). These Korean strings are CONTEXT ONLY for extracting visible-space cues. You MUST:
- IGNORE every proper name (character names, place names) in the Korean source — never reproduce them.
- Output the t2i prompt in English ASCII only. NEVER include Korean characters, Hanja, or kana.
- Use COMMON NOUNS only.

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.
- Chain 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 so the image model reuses them.
- For root anchors: thoroughly describe wall finish, floor, ceiling, lighting fixtures, palette — no parent photo to inherit from.
- Be concrete: name doors, windows, fixtures, furniture in frame.

Return strict JSON: {"t2i_prompt": "<single English string>"}.
"""

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


def collect_context(run_dir: Path, base_plan_ids: list[str],
                     extra_keywords: list[str] | None = None,
                     use_scope: bool = True) -> dict:
    """planning_gpt와 동일한 scope-based 전수조사.

    Scope detection 우선순위:
      1) context.scope.scenes (ground truth)
      2) shot_assignment에 base_plan_id가 in-scope인 shot들의 scene
      3) keyword fallback (위 둘이 비었을 때만)
    """
    bp_set = set(base_plan_ids)
    extra_keywords = extra_keywords or ["옥탑", "옥상", "rooftop"]
    ctx: dict = {"base_plans": [], "base_photos": [], "shots": [],
                 "spatial": None}

    sp2_path = run_dir / "step2_plan_specs.json"
    sp2 = json.loads(sp2_path.read_text(encoding="utf-8")) if sp2_path.exists() else {}
    for p in sp2.get("base_plans", []):
        if p.get("id") in bp_set:
            ctx["base_plans"].append(p)

    sp7_path = run_dir / "step7_photo_specs.json"
    if sp7_path.exists():
        d = json.loads(sp7_path.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)

    assign_map: dict = {}
    for a in sp2.get("shot_assignment", []):
        si = a.get("scene_index"); sx = a.get("shot_index")
        assign_map[(si, sx)] = (a.get("base_plan_id"), a.get("visual_domain"))

    ctx_path = run_dir / "context.json"
    ctx_json = json.loads(ctx_path.read_text(encoding="utf-8")) if ctx_path.exists() else {}

    scenes_in_scope: set = set()
    if use_scope:
        for si in (ctx_json.get("scope", {}).get("scenes") or []):
            scenes_in_scope.add(si)
    for (si, sx), (bp, _) in assign_map.items():
        if bp in bp_set and si is not None:
            scenes_in_scope.add(si)
    if not scenes_in_scope:
        for sc in ctx_json.get("scenes", []):
            h = sc.get("heading", "")
            txt = (sc.get("text") or "")
            if any(kw in h or kw in txt for kw in extra_keywords):
                scenes_in_scope.add(sc.get("scene_index"))

    seen: set = set()
    for s in ctx_json.get("shots", []):
        si = s.get("scene_index"); sx = s.get("shot_index")
        key = (si, sx)
        if key in seen: continue
        bp, dom = assign_map.get(key, (None, None))
        if si not in scenes_in_scope and bp not in bp_set:
            continue
        seen.add(key)
        step3_path = run_dir / f"step3_shot_S{si:02d}_Shot{sx}.json"
        step3 = json.loads(step3_path.read_text(encoding="utf-8")) if step3_path.exists() else None
        ctx["shots"].append({
            "scene_index": si,
            "shot_index": sx,
            "shot_id": f"S{si:02d}_Shot{sx}",
            "description": s.get("description") or "",
            "characters": s.get("characters") or [],
            "base_plan_id": bp,
            "visual_domain": dom,
            "step3": step3,
            "in_scope_reason": "shot_assignment" if (bp in bp_set) else (
                "context.scope.scenes" if si in (ctx_json.get("scope", {}).get("scenes") or []) else "scene_keyword"),
        })

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


def build_node_meta(plan: dict, ctx: dict) -> dict:
    """node_id → {kind, photo (kind=base_photo일 때), shots[] (node.shot_ids로 aggregate된 ctx shots)}.

    v4 schema에서는 모든 노드가 anchor_synthetic|base_photo이고 원본 shot은 node.shot_ids에 묶임.
    이 메타는 build_per_node_user_prompt에서 각 노드의 모든 원본 shot의 camera/additions/description을
    LLM에 전달하기 위해 사용.
    """
    # ctx shots를 shot_id로 인덱싱
    shots_by_id: dict = {}
    for s in ctx.get("shots", []):
        si = s.get("scene_index"); sx = s.get("shot_index")
        if si is None or sx is None: continue
        shots_by_id[f"S{si:02d}_Shot{sx}"] = s

    out: dict = {}
    for n in plan.get("nodes", []):
        nid = n["id"]; kind = n.get("kind"); src = n.get("source") or ""
        meta: dict = {"kind": kind, "photo": None, "shots": []}
        if kind == "base_photo":
            for ph in ctx.get("base_photos", []):
                if ph.get("id") == src or ph.get("id") == nid:
                    meta["photo"] = ph; break
        for sid in (n.get("shot_ids") or []):
            shot = shots_by_id.get(sid)
            if shot:
                meta["shots"].append(shot)
        out[nid] = meta
    return out


def find_anchor_plan_png(node: dict, plan: dict, ctx: dict, run_dir: Path) -> Path | None:
    """노드의 source_plan_id를 직접 사용. 없으면 base_photo의 source_plan_id, 그것도 없으면 None."""
    bp_id = (node.get("source_plan_id") or "").strip()
    # base_photo은 ctx.base_photos.source_plan_id로 검증/추론 가능
    if not bp_id and node.get("kind") == "base_photo":
        for ph in ctx.get("base_photos", []):
            if ph.get("id") in (node.get("source"), node.get("id")):
                bp_id = ph.get("source_plan_id") or ""
                break
    if not bp_id:
        return None
    p = run_dir / f"base_plan_{bp_id}.png"
    return p if p.exists() else 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. Reuse 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:
        if meta.get("kind") == "base_photo" and meta.get("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']}")

        shots = meta.get("shots") or []
        if shots:
            parts.append(
                f"\n== ORIGINAL SHOTS THIS BACKGROUND COVERS ({len(shots)} shots) =="
            )
            parts.append(
                "Each shot below specifies how the camera will frame this same background. "
                "Synthesize the prompt so the rendered image is wide/composed enough that "
                "ALL of these shots can be cropped/framed from it. The Korean description "
                "fields are CONTEXT ONLY — extract spatial/lighting/state cues but do not "
                "reproduce Korean proper names."
            )
            for s in shots:
                si = s.get("scene_index"); sx = s.get("shot_index")
                sid = f"S{si:02d}_Shot{sx}"
                parts.append(f"\n--- {sid} (base_plan={s.get('base_plan_id')}, "
                             f"domain={s.get('visual_domain')}, "
                             f"reason={s.get('in_scope_reason')}) ---")
                if s.get("description"):
                    parts.append(f"description (Korean source — ignore proper names): "
                                 f"{s['description']}")
                step3 = s.get("step3")
                if step3:
                    cam = step3.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: "
                                 f"{cam.get('height')} / {cam.get('fov')} / "
                                 f"{cam.get('lens_note')}")
                    for a in step3.get("additions") or []:
                        parts.append(f"  + [{a.get('type')}] @ {a.get('position', '')}: "
                                     f"{a.get('note', '')}")

    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_gpt_text_json(client: OpenAI, model: str, system: str, user: str,
                       schema: dict) -> dict:
    resp = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        response_format={
            "type": "json_schema",
            "json_schema": {
                "name": "T2IPromptOutput",
                "schema": schema,
                "strict": True,
            },
        },
    )
    txt = resp.choices[0].message.content or ""
    return json.loads(txt)


def gpt_image_edit_with_ref(client: OpenAI, model: str, ref_path: Path,
                            prompt: str, size: str, quality: str,
                            out_path: Path) -> Path:
    logger.info("[image edit] %s + ref[%s] → %s (%s, %s)",
                model, ref_path.name, out_path.name, size, quality)
    with open(ref_path, "rb") as f:
        resp = client.images.edit(
            model=model, image=[f], prompt=prompt,
            size=size, quality=quality, n=1,
        )
    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("--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="gpt-5.5")
    p.add_argument("--image-model", default="gpt-image-2")
    p.add_argument("--size", default="1024x1024")
    p.add_argument("--quality", default="high")
    p.add_argument("--prompts-only", action="store_true")
    p.add_argument("--skip-prompts", action="store_true")
    p.add_argument("--selected-shots", default="",
                   help="콤마 구분 selected shot ids (e.g. 'S05_Shot1,S10_Shot1,...'). "
                        "이 옵션 또는 --auto-shot-assignment 지정 시에만 selected 필터링 활성. "
                        "default는 모든 노드 image 생성.")
    p.add_argument("--auto-shot-assignment", action="store_true",
                   help="step2.shot_assignment에서 selected shot 자동 추출 + ancestor만 image 생성. "
                        "기본은 모든 노드 image 생성.")
    p.add_argument("--allow-text-only-roots", action="store_true",
                   help="root anchor의 plan PNG 없을 때 text-only generate fallback 허용 "
                        "(default는 fail-fast).")
    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("OPENAI_API_KEY"):
        logger.error("OPENAI_API_KEY not set"); return 1
    client = OpenAI()

    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: %s", nid); continue
            pid = n.get("parent_id") or ""
            parent = None
            if pid and pid not in ("null", "None"):
                parent = nodes_by_id.get(pid)
            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) %d chars",
                        i, len(execution_order), nid,
                        (parent or {}).get("id") or "(root)", len(user))
            try:
                resp = call_gpt_text_json(client, 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, {})
        pid = n.get("parent_id") or ""
        if pid in ("", "null", "None"): pid = "(root)"
        print(f"\n--- {nid}  [{n.get('kind')}, parent={pid}] ---")
        print(rendered_prompts.get(nid, "(missing)"))

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

    # === Step 1.5: render-set 결정 ===
    # default: 모든 노드 image 생성
    # --selected-shots OR --auto-shot-assignment: 명시한 shots만 + chain ancestor
    use_filter = bool(args.selected_shots.strip()) or args.auto_shot_assignment
    if not use_filter:
        render_set = set(execution_order)
        logger.info("default mode → render every node (%d)", len(render_set))
    else:
        selected_shot_ids: set[str] = set()
        if args.selected_shots.strip():
            selected_shot_ids = {s.strip() for s in args.selected_shots.split(",") if s.strip()}
            logger.info("--selected-shots provided: %d ids", len(selected_shot_ids))
        else:
            # --auto-shot-assignment
            sp2_path = run_dir / "step2_plan_specs.json"
            if sp2_path.exists():
                sp2 = json.loads(sp2_path.read_text(encoding="utf-8"))
                for a in sp2.get("shot_assignment", []):
                    si = a.get("scene_index"); sx = a.get("shot_index")
                    if si is not None and sx is not None:
                        selected_shot_ids.add(f"S{si:02d}_Shot{sx}")
                logger.info("auto-extracted from step2.shot_assignment: %d selected shots",
                            len(selected_shot_ids))

        # selected shot이 속한 노드 찾기
        directly_used: set[str] = set()
        for n in plan.get("nodes", []):
            if any(sid in selected_shot_ids for sid in (n.get("shot_ids") or [])):
                directly_used.add(n["id"])
        logger.info("nodes directly using selected shots: %d", len(directly_used))

        # 그 노드들의 모든 chain ancestor 추가 (image gen에 필요)
        render_set = set(directly_used)
        for nid in list(directly_used):
            cur = nid
            while True:
                n = nodes_by_id.get(cur)
                if not n: break
                pid = n.get("parent_id") or ""
                if not pid or pid in ("null", "None"): break
                if pid not in nodes_by_id: break
                render_set.add(pid)
                cur = pid
        logger.info("render-set after adding ancestors: %d nodes (skip %d)",
                    len(render_set), len(execution_order) - len(render_set))

        skipped = [n for n in execution_order if n not in render_set]
        if skipped:
            print("\n--- nodes SKIPPED for image gen (no selected shot, no chain need) ---")
            for nid in skipped:
                n = nodes_by_id.get(nid, {})
                print(f"  - {nid}  [{n.get('kind', '?')}, d={n.get('depth', '?')}, shots={len(n.get('shot_ids') or [])}]")
        print(f"\n--- nodes TO RENDER ({len(render_set)}) ---")
        for nid in execution_order:
            if nid in render_set:
                n = nodes_by_id.get(nid, {})
                src = "selected" if nid in directly_used else "chain ancestor"
                print(f"  + {nid}  [{n.get('kind', '?')}, d={n.get('depth', '?')}, shots={len(n.get('shot_ids') or [])}]  ← {src}")
        print()

    # === Step 2: 노드별 image chain 생성 (render_set만) ===
    rendered_images: dict[str, Path] = {}

    for i, nid in enumerate(execution_order, 1):
        if nid not in render_set:
            continue
        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", nid); continue

        out_p = out_dir / f"{nid}.png"
        pid = n.get("parent_id") or ""

        if pid and pid not in ("null", "None", ""):
            # 자식 노드 — 부모 image as ref
            ref_path = rendered_images.get(pid)
            if not ref_path:
                logger.error("parent image missing for %s (parent=%s) — skip", nid, pid); continue
            logger.info("[image %d/%d] node=%s with parent ref [%s]",
                        i, len(execution_order), nid, ref_path.name)
            try:
                gpt_image_edit_with_ref(client, args.image_model, ref_path,
                                         prompt_text, args.size, args.quality, out_p)
                rendered_images[nid] = out_p
            except Exception as e:
                logger.error("image gen failed for %s: %s", nid, e)
        else:
            # 루트 anchor — 도면 PNG as ref REQUIRED (Medium 5: fail-fast, no text-only fallback)
            plan_png = find_anchor_plan_png(n, plan, ctx, run_dir)
            if not plan_png:
                spid = n.get("source_plan_id") or "(empty)"
                logger.error(
                    "FAIL-FAST: root anchor %s has no floor plan PNG. "
                    "source_plan_id=%s. Either fix the plan's source_plan_id or "
                    "re-run with --allow-text-only-roots if you really want it.",
                    nid, spid)
                if not args.allow_text_only_roots:
                    return 1
                logger.warning("--allow-text-only-roots → falling back to text-only generate for %s", nid)
                try:
                    resp = client.images.generate(
                        model=args.image_model, prompt=prompt_text,
                        size=args.size, quality=args.quality, n=1,
                    )
                    out_p.write_bytes(base64.b64decode(resp.data[0].b64_json))
                    rendered_images[nid] = out_p
                    logger.info("  saved %s (%d KB)", out_p.name, out_p.stat().st_size // 1024)
                except Exception as e:
                    logger.error("image gen failed for %s: %s", nid, e)
            else:
                logger.info("[image %d/%d] node=%s with plan ref [%s]",
                            i, len(execution_order), nid, plan_png.name)
                try:
                    gpt_image_edit_with_ref(client, args.image_model, plan_png,
                                             prompt_text, args.size, args.quality, out_p)
                    rendered_images[nid] = out_p
                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())