"""[임시] chain_render_gpt_v6 배경 image를 ref로 selected shot의 실제 사진(인물 포함) 생성.

원본 chain_*.py 파일들에 영향 없음. 별도 임시 스크립트.

흐름:
  selected shots 11개 each (step2.shot_assignment 기준):
    1) chain v6 트리에서 매핑 노드 찾기 (node.shot_ids에 sid 포함)
    2) 그 노드의 배경 image PNG = chain_render_gpt_v6/{node_id}.png
    3) context.shots[i].description (한글 raw) + characters
    4) context.staging[shot_id]에서 camera_direction, lighting_mood 등 (있으면, 영문)
    5) ctx.characters에서 등장 인물의 description 추출
    6) LLM(gpt-5.5): 모든 정보 → shot 전용 t2i prompt (영문, common nouns)
    7) gpt-image-2 image edit: 배경 image as ref + 새 t2i → shot photo

체이닝 X — 각 shot은 자기 노드의 배경만 ref. 인물은 prompt 안 텍스트로만 등장.
"""
from __future__ import annotations

import argparse
import base64
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("temp_shot_render")


T2I_GEN_SYSTEM = """You write a single photorealistic image-generation prompt for ONE film shot.

Pipeline:
- The previously rendered BACKGROUND image will be supplied as the reference image.
- Your prompt MUST produce a photo that takes place INSIDE that exact background — same wall finish, floor, ceiling, lighting tone, color palette, furniture/door positions.
- The prompt adds the SHOT's narrative content: people, action, props, the camera framing.

INPUT LANGUAGE NOTICE:
- Shot description, character info, and dialogue are Korean source text. They are CONTEXT ONLY for extracting visual cues.
- Output prompt MUST be ASCII English only — NO Korean characters, NO Hanja, NO kana.
- Use COMMON NOUNS only. NEVER reproduce Korean proper names anywhere.
- For people: use generic descriptors based on the Korean character description, e.g. "a Korean woman in her early 20s, slim, mid-length black hair, casual student look with a backpack". DO NOT use the character's Korean name.
- For places: same — generic terms like "rooftop dwelling", "rear courtyard", never the location's Korean name.

Constraints in your output prompt:
- Open with: "Same room/space as the reference image — match its wall finish, floor, ceiling, lighting tone, color palette."
- Then describe the camera angle/lens/distance for THIS shot (use staging.camera_direction info if present, otherwise infer from description).
- Then describe the people/action: count, generic ethnicity/age/build, clothing, gesture/pose, position in frame.
- Then describe any narrative props/state changes specific to this shot.
- 35mm cinematic still aesthetic, photorealistic.
- NO blood, NO broken glass, NO gore, NO explicit violence. If the source describes injury, render only the aftermath atmosphere (worn surface, dust, shadow) — never graphic detail.

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

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


def build_shot_user_prompt(shot: dict, staging: dict | None,
                           characters: list[dict],
                           node_label: str, node_description: str) -> str:
    parts: list[str] = []
    parts.append("== BACKGROUND NODE THIS SHOT BELONGS TO ==")
    parts.append(f"node_label: {node_label}")
    parts.append(f"background_description: {node_description}")

    parts.append("\n== SHOT INFO ==")
    si = shot.get("scene_index"); sx = shot.get("shot_index")
    parts.append(f"id: S{si:02d}_Shot{sx}")
    parts.append(f"description (Korean source — extract visual cues, ignore proper names):")
    parts.append(f"  {shot.get('description', '')}")

    chars_in_shot = shot.get("characters") or []
    if chars_in_shot:
        parts.append(f"\n== CHARACTERS IN THIS SHOT ({len(chars_in_shot)}) ==")
        parts.append("Use the descriptions below to render generic English descriptors. NEVER use the Korean names.")
        for cn in chars_in_shot:
            for c in characters:
                if c.get("name") == cn:
                    parts.append(f"\n--- {c.get('short_id')} ---")
                    parts.append(f"description (Korean source): {c.get('description', '')}")
                    break

    if staging and (staging.get("camera_direction") or staging.get("lighting_mood")):
        parts.append("\n== STAGING (English source — use directly) ==")
        if staging.get("camera_direction"):
            parts.append(f"camera_direction: {staging['camera_direction']}")
        if staging.get("lighting_mood"):
            parts.append(f"lighting_mood: {staging['lighting_mood']}")
        if staging.get("perspective"):
            parts.append(f"perspective: {staging['perspective']}")
        kbe = staging.get("key_bg_elements") or []
        if kbe:
            parts.append(f"key_bg_elements: {json.dumps(kbe, ensure_ascii=False)}")

    parts.append("\n== TASK ==")
    parts.append("Write ONE photorealistic t2i prompt for this shot, set inside the background node above.")
    parts.append('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("--planning-dir", required=True,
                   help="chain_planning_gpt_v6 (chain_structure.json 위치)")
    p.add_argument("--background-dir", required=True,
                   help="chain_render_gpt_v6 (배경 image PNG들 위치)")
    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="콤마 구분 명시 (비워두면 step2.shot_assignment 자동)")
    args = p.parse_args()

    run_dir = Path(args.run_dir).resolve()
    planning_dir = Path(args.planning_dir).resolve()
    bg_dir = Path(args.background_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()

    ctx = json.loads((run_dir / "context.json").read_text(encoding="utf-8"))
    plan = json.loads((planning_dir / "chain_structure.json").read_text(encoding="utf-8"))
    sp2 = json.loads((run_dir / "step2_plan_specs.json").read_text(encoding="utf-8"))

    # selected shots
    if args.selected_shots.strip():
        sids = [s.strip() for s in args.selected_shots.split(",") if s.strip()]
        selected = []
        for sid in sids:
            # parse "S04_Shot1" → (4, 1, "S04_Shot1")
            try:
                si = int(sid.split("_")[0][1:]); sx = int(sid.split("Shot")[1])
                selected.append((si, sx, sid))
            except Exception:
                logger.warning("invalid shot id format: %s", sid)
    else:
        selected = []
        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.append((si, sx, f"S{si:02d}_Shot{sx}"))
    logger.info("selected shots: %d", len(selected))

    # shot_id → chain node 매핑 (first-match wins, 중복은 warning)
    node_for_shot: dict = {}
    for n in plan.get("nodes", []):
        for sid in (n.get("shot_ids") or []):
            if sid in node_for_shot:
                logger.warning(
                    "shot %s appears in multiple nodes — keeping first %s, ignoring %s",
                    sid, node_for_shot[sid]["id"], n["id"])
                continue
            node_for_shot[sid] = n

    shot_by_id = {f"S{s['scene_index']:02d}_Shot{s['shot_index']}": s for s in ctx.get("shots", [])}
    staging_by_id = {f"S{sg['scene_index']:02d}_Shot{sg['shot_index']}": sg for sg in ctx.get("staging", [])}
    characters = ctx.get("characters", [])

    # === Step 1: 노드별 t2i prompt 생성 ===
    prompts_p = out_dir / "shot_prompts.json"
    if args.skip_prompts and prompts_p.exists():
        rendered_prompts = json.loads(prompts_p.read_text(encoding="utf-8"))
        logger.info("[skip prompts] reusing shot_prompts.json (%d)", len(rendered_prompts))
    else:
        rendered_prompts = {}
        for i, (si, sx, sid) in enumerate(selected, 1):
            shot = shot_by_id.get(sid, {})
            staging = staging_by_id.get(sid)
            node = node_for_shot.get(sid)
            if not node:
                logger.error("no chain node found for %s — skip", sid); continue
            user_prompt = build_shot_user_prompt(
                shot, staging, characters,
                node.get("label", ""), node.get("description", ""),
            )
            logger.info("[prompt %d/%d] %s → node=%s (%d chars)",
                        i, len(selected), sid, node["id"], len(user_prompt))
            try:
                resp = call_gpt_text_json(client, args.text_model,
                                           T2I_GEN_SYSTEM, user_prompt, T2I_GEN_SCHEMA)
                rendered_prompts[sid] = {
                    "shot_id": sid,
                    "node_id": node["id"],
                    "node_label": node.get("label"),
                    "prompt": resp.get("t2i_prompt", ""),
                }
            except Exception as e:
                logger.error("prompt gen failed for %s: %s", sid, e)
        prompts_p.write_text(json.dumps(rendered_prompts, ensure_ascii=False, indent=2),
                             encoding="utf-8")
        logger.info("saved shot_prompts.json")

    # Print for review
    print("\n" + "=" * 70)
    print("=== PER-SHOT T2I PROMPTS ===")
    print("=" * 70)
    for sid, info in rendered_prompts.items():
        print(f"\n--- {sid}  → node={info['node_id']} ({info.get('node_label', '')}) ---")
        print(info.get("prompt", "(missing)"))

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

    # === Step 2: 배경 ref로 shot image 생성 ===
    for i, (si, sx, sid) in enumerate(selected, 1):
        info = rendered_prompts.get(sid)
        if not info:
            continue
        node_id = info["node_id"]
        bg_path = bg_dir / f"{node_id}.png"
        if not bg_path.exists():
            logger.error("background image missing for %s: %s — skip", sid, bg_path); continue
        out_p = out_dir / f"shot_{sid}.png"
        logger.info("[image %d/%d] %s using bg=%s",
                    i, len(selected), sid, bg_path.name)
        try:
            gpt_image_edit_with_ref(client, args.image_model, bg_path,
                                     info["prompt"], args.size, args.quality, out_p)
        except Exception as e:
            logger.error("image gen failed for %s: %s", sid, e)

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


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