"""기존 v4 run 디렉토리의 gallery.html을 prompt 포함해서 재생성.

사용법:
  python scripts/regen_gallery_with_prompts.py <run-dir>
"""
from __future__ import annotations

import argparse
import html
import json
import sys
from pathlib import Path
from typing import Dict, Any, List, Optional


def load_json(p: Path) -> Dict[str, Any]:
    if not p.exists():
        return {}
    return json.loads(p.read_text(encoding="utf-8"))


def find_first(run_dir: Path, prefix: str, plan_id: str) -> Optional[str]:
    """파일명 매칭 — exact prefix_<plan_id>.png."""
    for p in run_dir.iterdir():
        if p.is_file() and p.name == f"{prefix}_{plan_id}.png":
            return p.name
    return None


def collect_rows(run_dir: Path) -> List[Dict[str, Any]]:
    plans = load_json(run_dir / "step2_plan_specs.json")
    photos = load_json(run_dir / "step7_photo_specs.json")
    shot_files = sorted(run_dir.glob("step3_shot_*.json"))
    shots = [load_json(p) for p in shot_files]

    photos_dir = run_dir / "photos"
    rows: List[Dict[str, Any]] = []

    # base 행 — base_plan + isometric + axo + base_photo
    plan_id_to_base_photo_t2i: Dict[str, str] = {}
    for ph in photos.get("base_photos", []):
        spid = ph.get("source_plan_id")
        if spid:
            plan_id_to_base_photo_t2i[spid] = ph.get("t2i_prompt", "")

    for p in plans.get("base_plans", []):
        pid = p["id"]
        kind = "ANCHOR" if p.get("is_anchor") else "BASE"
        base_plan_file = find_first(run_dir, "base_plan", pid)
        iso_file = find_first(run_dir, "isometric", pid)
        axo_file = find_first(run_dir, "axo", pid)
        # photo 파일명 — pid 그대로 + photo_base_ prefix
        photo_file = None
        if (photos_dir / f"photo_base_{pid}.png").exists():
            photo_file = f"photo_base_{pid}.png"
        rows.append({
            "label_id": pid,
            "label_ko": p.get("label", ""),
            "domain": p.get("visual_domain", ""),
            "kind": kind,
            "scale": p.get("scale_estimate", ""),
            "is_shot": False,
            "plan_file": base_plan_file,
            "iso_file": iso_file,
            "axo_file": axo_file,
            "photo_file": photo_file,
            "plan_t2i": p.get("t2i_prompt", ""),
            "photo_t2i": plan_id_to_base_photo_t2i.get(pid, ""),
        })

    # shot 행
    shot_id_to_photo_t2i: Dict[str, str] = {}
    for ph in photos.get("shot_photos", []):
        src = ph.get("source_shot") or {}
        si = src.get("scene_index"); sx = src.get("shot_index")
        if si is not None and sx is not None:
            try:
                lbl = f"S{int(si):02d}_Shot{int(sx)}"
                shot_id_to_photo_t2i[lbl] = ph.get("t2i_prompt", "")
            except Exception:
                pass

    for sh in shots:
        si = sh.get("scene_index"); sx = sh.get("shot_index")
        if si is None or sx is None:
            continue
        label = f"S{int(si):02d}_Shot{int(sx)}"
        shot_plan_file = f"shot_plan_{label}.png" if (run_dir / f"shot_plan_{label}.png").exists() else None
        shot_photo_file = f"photo_shot_{label}.png" if (photos_dir / f"photo_shot_{label}.png").exists() else None
        rows.append({
            "label_id": label,
            "label_ko": "",
            "domain": sh.get("visual_domain", ""),
            "kind": "SHOT",
            "scale": "",
            "is_shot": True,
            "plan_file": shot_plan_file,
            "iso_file": None,
            "axo_file": None,
            "photo_file": shot_photo_file,
            "plan_t2i": sh.get("t2i_prompt", ""),
            "photo_t2i": shot_id_to_photo_t2i.get(label, ""),
        })

    return rows


def render_html(run_dir: Path, rows: List[Dict[str, Any]]) -> str:
    manifest = load_json(run_dir / "manifest.json")
    val = manifest.get("validation") or {}
    badges = "".join(
        f'<span class="badge {"ok" if not val.get(k) else "fail"}">{k.replace("_", " ")} {len(val.get(k) or [])}</span>'
        for k in ("scenario_word_hits", "unsafe_word_hits", "missing_files",
                  "schema_errors", "reference_integrity_errors")
    )

    def img_or_blank(fname: Optional[str], subdir: str = "") -> str:
        if not fname:
            return '<em class="missing">missing</em>'
        src = f"{subdir}{fname}" if not subdir or fname.startswith(subdir) else f"{subdir}{fname}"
        return f'<img src="{html.escape(src)}" loading="lazy">'

    def prompt_block(label: str, text: str) -> str:
        if not text:
            return ""
        return (
            f'<details class="prompt"><summary>{html.escape(label)} t2i_prompt '
            f'({len(text)} chars)</summary>'
            f'<pre>{html.escape(text)}</pre></details>'
        )

    body_rows: List[str] = []
    for r in rows:
        kind_cls = r["kind"].lower()
        # 도면 panes: base_plan / iso / axo / photo
        panes = []
        panes.append(("PLAN", img_or_blank(r["plan_file"])))
        if r["iso_file"]:
            panes.append(("ISO", img_or_blank(r["iso_file"])))
        if r["axo_file"]:
            panes.append(("AXO", img_or_blank(r["axo_file"])))
        panes.append(("PHOTO", img_or_blank(r["photo_file"], "photos/")))

        panes_html = "".join(
            f'<figure><div class="imgwrap">{img_html}</div>'
            f'<figcaption><span class="kind-tag kind-{kind.lower()}">{kind}</span></figcaption></figure>'
            for kind, img_html in panes
        )

        prompts_html = (
            prompt_block("PLAN", r.get("plan_t2i", ""))
            + prompt_block("PHOTO", r.get("photo_t2i", ""))
        )

        body_rows.append(f"""
<div class="row" data-kind="{kind_cls}">
  <div class="label-row">
    <span class="id">{html.escape(r["label_id"])}</span>
    <span class="kind kind-{kind_cls}">{r["kind"]}</span>
    <span class="domain">{html.escape(r["domain"])}</span>
    <span class="scale">{html.escape(r["scale"])}</span>
    <span class="ko">{html.escape(r["label_ko"])}</span>
  </div>
  <div class="panes panes-{len(panes)}">{panes_html}</div>
  <div class="prompts">{prompts_html}</div>
</div>
""")

    return f"""<!doctype html>
<html lang="ko"><head><meta charset="utf-8">
<title>v4 Gallery (with prompts) — {manifest.get("run_id","")}</title>
<style>
  body {{ margin:0;background:#0e0f12;color:#e8eaed;font-family:-apple-system,sans-serif;line-height:1.5;}}
  header {{ padding:18px 28px;border-bottom:1px solid #2a2d34;background:#14161b;position:sticky;top:0;z-index:10;}}
  header h1 {{ margin:0 0 6px;font-size:18px;}}
  .meta {{ color:#9aa0a6;font-size:12px;}}
  .badge {{ display:inline-block;padding:2px 9px;margin-right:5px;border-radius:10px;font-size:10px;}}
  .badge.ok {{ background:#15331c;color:#74e896;border:1px solid #1f5230;}}
  .badge.fail {{ background:#3a1c1c;color:#ff9090;border:1px solid #5a2828;}}
  main {{ padding:20px;max-width:2000px;margin:0 auto;}}
  .row {{ background:#181a1f;border:1px solid #2a2d34;border-radius:8px;margin-bottom:18px;padding:12px;}}
  .label-row {{ display:flex;gap:12px;align-items:baseline;padding:6px 10px;background:#14161b;border-left:3px solid #7aa2f7;border-radius:6px;margin-bottom:8px;}}
  .id {{ font-family:ui-monospace,monospace;color:#7aa2f7;font-size:13px;}}
  .domain {{ font-size:11px;color:#9aa0a6;background:#222;padding:1px 8px;border-radius:8px;}}
  .scale {{ font-size:11px;color:#9aa0a6;}}
  .ko {{ font-size:12px;color:#cfd2d7;}}
  .kind {{ font-size:10px;padding:1px 7px;border-radius:8px;}}
  .kind-anchor, .kind-tag.kind-anchor {{ background:#3a3a52;color:#c9b3ff;}}
  .kind-base, .kind-tag.kind-base {{ background:#1f3a52;color:#9ec5fe;}}
  .kind-shot, .kind-tag.kind-shot {{ background:#3a3320;color:#ffd47a;}}
  .kind-tag {{ font-family:ui-monospace,monospace;font-size:10px;padding:1px 7px;border-radius:8px;}}
  .panes {{ display:grid;gap:10px;}}
  .panes-2 {{ grid-template-columns: repeat(2, 1fr);}}
  .panes-3 {{ grid-template-columns: repeat(3, 1fr);}}
  .panes-4 {{ grid-template-columns: repeat(4, 1fr);}}
  figure {{ margin:0;background:#fff;border-radius:6px;overflow:hidden;}}
  .imgwrap {{ height:280px;display:flex;align-items:center;justify-content:center;}}
  img {{ max-width:100%;max-height:100%;object-fit:contain;}}
  figcaption {{ padding:4px 8px;background:#1d2026;font-size:10px;color:#9aa0a6;border-top:1px solid #2a2d34;text-align:center;}}
  .missing {{ color:#a06060;font-style:italic;font-size:11px;}}
  .prompts {{ margin-top:10px;display:flex;flex-direction:column;gap:6px;}}
  details.prompt {{ background:#0e0f12;border:1px solid #2a2d34;border-radius:6px;padding:0;}}
  details.prompt summary {{ cursor:pointer;padding:6px 10px;font-size:12px;color:#9aa0a6;font-weight:500;}}
  details.prompt summary:hover {{ color:#e8eaed;}}
  details.prompt[open] summary {{ color:#7aa2f7;border-bottom:1px solid #2a2d34;}}
  details.prompt pre {{ margin:0;padding:10px 12px;font-size:11px;color:#cfd2d7;white-space:pre-wrap;word-break:break-word;font-family:ui-monospace,monospace;line-height:1.45;}}
</style></head><body>
<header>
  <h1>v4 Gallery (with prompts) — Plan ↔ Photo</h1>
  <div class="meta">
    <code>{manifest.get("run_id","")}</code> ·
    text=<code>{manifest.get("models",{}).get("text","?")}</code> ·
    image=<code>{manifest.get("models",{}).get("image","?")}</code><br>
    {badges}
  </div>
</header>
<main>
{"".join(body_rows)}
</main>
</body></html>"""


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("run_dir", type=Path)
    p.add_argument("--out", default=None, help="기본: <run_dir>/gallery.html (덮어쓰기)")
    args = p.parse_args()
    if not args.run_dir.is_dir():
        print(f"run-dir 없음: {args.run_dir}", file=sys.stderr); return 1
    rows = collect_rows(args.run_dir)
    out_path = Path(args.out) if args.out else (args.run_dir / "gallery.html")
    out_path.write_text(render_html(args.run_dir, rows), encoding="utf-8")
    print(f"saved: {out_path}  ({len(rows)} rows)")
    return 0


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