"""forest_exp 공용 라이브러리 — '전체 숲' 총괄 파이프라인 실험 (2026-07-04).

실험 전용(scratchpad, 커밋 금지). production 코드 차용(import) + 실데이터 읽기 전용.
- capture scope 미개방 → image_asset/체크포인트 write 0.
- 시나리오 중립: 이 파일의 어떤 템플릿에도 장소/인물 고유명사 없음.
  구체 내용은 전부 cp 데이터 문자열로 주입된다.
"""
from __future__ import annotations

import base64
import json
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

ROOT = Path("/Users/manta/Documents/Projects/TheRoad-I1")
BACKEND = ROOT / "backend"
EXP = ROOT / "scratchpad" / "forest_exp"
OUT = EXP / "out"
PLANS = EXP / "plans"
RUNLOG = EXP / "runlog.jsonl"

# 실데이터 (E2E) — 데이터 선택 파라미터일 뿐 코드 로직은 프로젝트 불문 동작.
PROJECT_ID = "b0ad5c18-5140-4022-a2b0-805de0e5a385"
EPISODE_ID = "0d30302a-51c8-4ac9-98e2-6cd9180512bf"
CP = ROOT / "projects" / PROJECT_ID / "checkpoints" / "episodes" / EPISODE_ID
EP_IMAGES = ROOT / "projects" / PROJECT_ID / "episodes" / EPISODE_ID / "images"

# backend 모듈 import 준비 (.env 는 cwd 기준 로드 → backend 로 chdir)
os.chdir(BACKEND)
if str(BACKEND) not in sys.path:
    sys.path.insert(0, str(BACKEND))

from app.core.config import settings  # noqa: E402
from app.modules.llm.llm_client import call_structured  # noqa: E402
from app.modules.llm.gpt_image_primitive import call_gpt_image_bytes  # noqa: E402
from app.modules.llm.gemini_key_pool import get_next_key  # noqa: E402
from app.modules.llm.safety import sanitize_for_safety  # noqa: E402
from app.modules.gemini_i2i_editor import _gemini_generate_content  # noqa: E402

_openai_client = None


def openai_client():
    global _openai_client
    if _openai_client is None:
        from openai import OpenAI
        _openai_client = OpenAI(
            api_key=settings.openai_api_key,
            timeout=float(settings.llm_timeout_image_gen),
        )
    return _openai_client


# ── 로깅 ──

def runlog(event: Dict[str, Any]) -> None:
    event = {"ts": time.strftime("%Y-%m-%d %H:%M:%S"), **event}
    RUNLOG.parent.mkdir(parents=True, exist_ok=True)
    with RUNLOG.open("a") as f:
        f.write(json.dumps(event, ensure_ascii=False) + "\n")


# ── 데이터 로더 ──

def load_step(step: str) -> Dict[str, Any]:
    return json.loads((CP / step / "manifest.json").read_text())["data"]


def pick_group(override: Optional[str] = None) -> Dict[str, Any]:
    """데이터 기준 그룹 선택: 실내+실외 멤버 공존(mixed) 그룹 중 총 shot_count 최대.

    is_indoor 가 bool 이 아닌 멤버는 제외(구조 계약 — 추측 금지).
    """
    groups = load_step("background_classify")["building_groups"]
    cands = []
    for g in groups:
        members = [m for m in (g.get("members") or [])
                   if isinstance(m.get("is_indoor"), bool)]
        if not members or not g.get("group_id"):
            continue
        has_in = any(m["is_indoor"] for m in members)
        has_out = any(not m["is_indoor"] for m in members)
        total = sum(int(m.get("shot_count") or 0) for m in members)
        if override:
            if g["group_id"] == override:
                return {**g, "members": members, "_total_shots": total}
            continue
        if has_in and has_out:
            cands.append(({**g, "members": members, "_total_shots": total}, total))
    if override:
        raise SystemExit(f"group not found: {override}")
    if not cands:
        raise SystemExit("no mixed indoor/outdoor group in building_groups")
    cands.sort(key=lambda x: -x[1])
    return cands[0][0]


def shot_key(scene_index: int, shot_index: int) -> str:
    return f"S{scene_index}_Shot{shot_index}"


def build_recon(group_override: Optional[str] = None) -> Dict[str, Any]:
    """그룹 → loc → bg → shot 조인 + fp/passport/베이스라인 자산 인벤토리."""
    group = pick_group(group_override)
    locs = {m["loc_id"]: m for m in group["members"]}

    mp = load_step("background_master_plan")
    catalog = mp["background_catalog"]
    group_bgs = {bid: e for bid, e in catalog.items() if e.get("loc_id") in locs}

    # shot 키 수집 (catalog applies_to_shots 기준 — bg 에 묶인 샷만이 배경 실험 대상)
    shot_to_bgs: Dict[str, List[str]] = {}
    for bid, e in group_bgs.items():
        for sk in e.get("applies_to_shots") or []:
            shot_to_bgs.setdefault(sk, []).append(bid)

    # shot 내용 조인
    se = load_step("shot_extract")["scenes"]
    desc_by_key: Dict[str, Dict[str, Any]] = {}
    for sc in se:
        for sh in sc.get("shots") or []:
            desc_by_key[shot_key(sc["scene_index"], sh["shot_index"])] = {
                "scene_heading": sc.get("scene_heading"),
                "description": sh.get("description"),
                "characters": sh.get("characters") or [],
            }
    staging_by_key = {
        shot_key(s["scene_index"], s["shot_index"]): s
        for s in load_step("shot_staging")["shots"]
    }
    detail_by_key: Dict[str, Dict[str, Any]] = {}
    for sc in load_step("scene_detail")["scenes"]:
        k = shot_key(sc["scene_index"], sc.get("_shot_index"))
        detail_by_key[k] = {
            "t2i_variations": sc.get("t2i_variations"),
            "visible_entities": sc.get("visible_entities"),
            "render_prompt_card": sc.get("render_prompt_card"),
        }

    shots: Dict[str, Dict[str, Any]] = {}
    for sk, bids in sorted(shot_to_bgs.items()):
        loc_ids = sorted({group_bgs[b]["loc_id"] for b in bids})
        is_indoor = all(locs[l]["is_indoor"] for l in loc_ids)
        si, shi = sk.split("_Shot")
        shots[sk] = {
            "scene_index": int(si[1:]),
            "shot_index": int(shi),
            "loc_ids": loc_ids,
            "is_indoor": is_indoor,
            "bg_ids": sorted(bids),
            **(desc_by_key.get(sk) or {}),
            "staging": staging_by_key.get(sk),
            "detail": detail_by_key.get(sk),
        }

    # fp 경로 (loc → catalog depends_on_fp → floor_plan_render png)
    fp_render = load_step("floor_plan_render")["floor_plans"]
    fp_by_loc: Dict[str, List[Dict[str, str]]] = {}
    for bid, e in group_bgs.items():
        for fp_id in e.get("depends_on_fp") or []:
            r = fp_render.get(fp_id) or {}
            if r.get("status") == "ok" and r.get("png_path"):
                lst = fp_by_loc.setdefault(e["loc_id"], [])
                if not any(x["fp_id"] == fp_id for x in lst):
                    lst.append({"fp_id": fp_id, "png_path": r["png_path"]})

    # 캐릭터 passport (DB read-only) + 이름→short_id 매핑
    char_ids = sorted({c for s in shots.values() for c in (s.get("characters") or [])
                       if isinstance(c, str)})
    passports = query_passports()
    name_to_sid = query_character_names()

    baseline_aerial = EP_IMAGES / "background_chain" / f"aerial_{group['group_id']}.png"

    recon = {
        "group_id": group["group_id"],
        "anchor_loc": group.get("anchor_loc"),
        "members": group["members"],
        "bgs": group_bgs,
        "shots": shots,
        "fp_by_loc": fp_by_loc,
        "characters_in_shots": char_ids,
        "passports": passports,
        "character_name_to_sid": name_to_sid,
        "baseline_aerial": str(baseline_aerial) if baseline_aerial.exists() else None,
    }
    return recon


def query_passports() -> Dict[str, str]:
    """캐릭터 short_id → primary reference PNG 절대경로 (DB read-only)."""
    sql = (
        "SELECT ec.short_id, ia.file_path FROM image_asset ia "
        "JOIN entity_canon ec ON ec.id::text=ia.entity_id::text "
        f"WHERE ia.project_id='{PROJECT_ID}' AND ia.asset_type='reference' "
        "AND ec.entity_type='character' AND ia.is_primary=1 AND ia.status='generated'"
    )
    env = {**os.environ, "PGPASSWORD": "theroad_dev_2026"}
    out = subprocess.run(
        ["psql", "-h", "localhost", "-U", "theroad", "-d", "theroad", "-tAc", sql],
        capture_output=True, text=True, env=env, check=True,
    ).stdout
    result: Dict[str, str] = {}
    for line in out.strip().splitlines():
        if "|" not in line:
            continue
        sid, path = line.split("|", 1)
        p = Path(path)
        if not p.is_absolute():
            p = ROOT / path
        if p.exists():
            result[sid.strip()] = str(p)
    return result


def query_character_names() -> Dict[str, str]:
    """캐릭터 이름 → short_id (DB read-only)."""
    sql = (
        "SELECT name, short_id FROM entity_canon "
        f"WHERE project_id='{PROJECT_ID}' AND entity_type='character'"
    )
    env = {**os.environ, "PGPASSWORD": "theroad_dev_2026"}
    out = subprocess.run(
        ["psql", "-h", "localhost", "-U", "theroad", "-d", "theroad", "-tAc", sql],
        capture_output=True, text=True, env=env, check=True,
    ).stdout
    result: Dict[str, str] = {}
    for line in out.strip().splitlines():
        if "|" in line:
            name, sid = line.rsplit("|", 1)
            result[name.strip()] = sid.strip()
    return result


def query_production_stills(scene_index: int, shot_index: int) -> List[str]:
    """해당 (scene, shot)의 production 최종 스틸 경로들 (베이스라인 육안 비교용)."""
    sql = (
        "SELECT ia.file_path FROM image_asset ia JOIN scene_still ss "
        "ON ss.id::text=ia.still_id::text "
        f"WHERE ia.project_id='{PROJECT_ID}' AND ia.asset_type='scene' "
        f"AND ss.scene_index={scene_index} AND ss.shot_index={shot_index} "
        "ORDER BY ia.created_at DESC"
    )
    env = {**os.environ, "PGPASSWORD": "theroad_dev_2026"}
    out = subprocess.run(
        ["psql", "-h", "localhost", "-U", "theroad", "-d", "theroad", "-tAc", sql],
        capture_output=True, text=True, env=env, check=True,
    ).stdout
    paths = []
    for line in out.strip().splitlines():
        p = Path(line.strip())
        if not p.is_absolute():
            p = ROOT / line.strip()
        if p.exists():
            paths.append(str(p))
    return paths


# ── LLM / 이미지 래퍼 ──

def llm(step_tag: str, system_prompt: str, user_prompt: Any,
        response_schema: Dict[str, Any], model: str = "gpt",
        temperature: float = 0.2) -> Dict[str, Any]:
    """call_structured 차용 — model 강제(project_config), 전체 원문 전달(자르기 금지)."""
    t0 = time.monotonic()
    result = call_structured(
        step_tag, system_prompt, user_prompt, response_schema,
        project_config={step_tag: {"model": model}},
        schema_name=step_tag, temperature=temperature,
    )
    user_text = user_prompt if isinstance(user_prompt, str) else json.dumps(
        [p for p in user_prompt if p.get("type") == "text"], ensure_ascii=False)
    runlog({"kind": "llm", "step": step_tag, "model": model,
            "dur_s": round(time.monotonic() - t0, 1),
            "system_prompt": system_prompt, "user_prompt": user_text})
    return result


def png_data_url(path: str | Path) -> Dict[str, Any]:
    b64 = base64.b64encode(Path(path).read_bytes()).decode("ascii")
    return {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}}


def img_gpt(tag: str, prompt: str, refs: Optional[Sequence[str | Path]] = None,
            size: str = "1536x864", quality: str = "high",
            out_path: Optional[Path] = None, max_attempts: int = 2) -> Path:
    """gpt-image-2 generate/edit (production call_gpt_image_bytes 차용).

    capture scope 미개방 상태로 호출 → DB 기록 없음. moderation 류 실패 시
    sanitize_for_safety 로 1회 재시도(bounded).
    """
    assert out_path is not None
    out_path.parent.mkdir(parents=True, exist_ok=True)
    if out_path.exists():
        runlog({"kind": "img", "tag": tag, "skip": "exists", "out": str(out_path)})
        return out_path
    mode = "edit" if refs else "generate"
    current = prompt
    t0 = time.monotonic()
    last: Optional[Exception] = None
    for attempt in range(1, max_attempts + 1):
        try:
            png = call_gpt_image_bytes(
                openai_client(), mode=mode, prompt=current,
                ref_paths=[str(r) for r in refs] if refs else None,
                call_kwargs={"model": "gpt-image-2", "size": size,
                             "quality": quality, "n": 1},
                edit_image_as_list=True,
            )
            out_path.write_bytes(png)
            runlog({"kind": "img", "engine": "gpt-image-2", "mode": mode, "tag": tag,
                    "refs": [str(r) for r in (refs or [])], "size": size,
                    "attempt": attempt, "dur_s": round(time.monotonic() - t0, 1),
                    "bytes": len(png), "out": str(out_path), "prompt": current})
            return out_path
        except Exception as exc:  # bounded retry (moderation 등)
            last = exc
            runlog({"kind": "img_error", "tag": tag, "attempt": attempt,
                    "error": str(exc)[:500]})
            current = sanitize_for_safety(current)
    raise RuntimeError(f"img_gpt failed for {tag}: {last}")


def img_nb2(tag: str, prompt: str, labeled_refs: Sequence[Tuple[str, str | Path]],
            aspect_ratio: str = "16:9", out_path: Optional[Path] = None,
            max_attempts: int = 2) -> Path:
    """nb2(gemini image) i2i 차용 — 라벨드 ref 멀티 입력, 최종/임시샷용."""
    assert out_path is not None
    out_path.parent.mkdir(parents=True, exist_ok=True)
    if out_path.exists():
        runlog({"kind": "img", "tag": tag, "skip": "exists", "out": str(out_path)})
        return out_path
    inputs = [(label, Path(p).read_bytes()) for label, p in labeled_refs]
    current = prompt
    t0 = time.monotonic()
    last: Optional[Exception] = None
    for attempt in range(1, max_attempts + 1):
        try:
            png = _gemini_generate_content(
                get_next_key(), settings.gemini_image_model, current, inputs,
                aspect_ratio=aspect_ratio,
            )
            out_path.write_bytes(png)
            runlog({"kind": "img", "engine": "nb2", "mode": "i2i", "tag": tag,
                    "refs": [str(p) for _, p in labeled_refs],
                    "ref_labels": [l for l, _ in labeled_refs],
                    "attempt": attempt, "dur_s": round(time.monotonic() - t0, 1),
                    "bytes": len(png), "out": str(out_path), "prompt": current})
            return out_path
        except Exception as exc:
            last = exc
            runlog({"kind": "img_error", "tag": tag, "attempt": attempt,
                    "error": str(exc)[:500]})
            current = sanitize_for_safety(current)
    raise RuntimeError(f"img_nb2 failed for {tag}: {last}")


# ── 공용 직렬화 헬퍼 (LLM 입력 — 원문 전체, 자르기 금지) ──

def shot_block(sk: str, s: Dict[str, Any], include_staging: bool = True) -> str:
    """한 샷의 내용을 LLM 입력 텍스트 블록으로 (전체 원문)."""
    lines = [f"[{sk}] ({'indoor' if s['is_indoor'] else 'outdoor'}; "
             f"loc={','.join(s['loc_ids'])}; scene_heading={s.get('scene_heading')})"]
    lines.append(f"description: {s.get('description')}")
    if s.get("characters"):
        lines.append(f"characters: {', '.join(s['characters'])}")
    st = s.get("staging") or {}
    if include_staging and st:
        lines.append(f"framing_scale: {st.get('framing_scale')}")
        lines.append(f"camera_direction: {st.get('camera_direction')}")
        kbe = st.get("key_bg_elements") or []
        if kbe:
            lines.append("key_bg_elements: " + json.dumps(kbe, ensure_ascii=False))
        fsc = st.get("frame_spatial_contract") or {}
        if fsc:
            lines.append("frame_spatial_contract: "
                         + json.dumps(fsc, ensure_ascii=False))
    return "\n".join(lines)


def members_block(members: List[Dict[str, Any]]) -> str:
    lines = []
    for m in members:
        tag = "indoor" if m.get("is_indoor") else "outdoor"
        lines.append(f"- ({tag}, {m.get('loc_id')}) {m.get('label')}: {m.get('summary')}")
    return "\n".join(lines)


def load_recon() -> Dict[str, Any]:
    return json.loads((PLANS / "recon.json").read_text())


def load_plan(name: str) -> Dict[str, Any]:
    return json.loads((PLANS / f"{name}.json").read_text())


def save_plan(name: str, data: Dict[str, Any]) -> Path:
    PLANS.mkdir(parents=True, exist_ok=True)
    p = PLANS / f"{name}.json"
    p.write_text(json.dumps(data, ensure_ascii=False, indent=1))
    return p