"""s1 — 뿌리(root): LLM 요소 선정(R1) → 통합 조감 스케치 생성(R2, A/B) → readback.

사용: .venv/bin/python s1_root.py [--only plan|image|readback]
산출: plans/root_plan.json, out/root/root_A.png(순수 T2I), root_B.png(배치도 I2I),
      plans/root_readback_{A,B}.json
"""
import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
import forest_lib as F  # noqa: E402

OUTDIR = F.OUT / "root"

R1_SYSTEM = """You are the master planner for a group of connected filming locations
(one real-world place made of indoor and outdoor parts). You are given the place's
member descriptions and the FULL content of every shot that happens there.

Your job: design ONE comprehensive TOP-DOWN aerial site sketch that lets every later
image of any of these shots stay spatially consistent. Select the elements that must
appear on that single sketch.

Hard rules:
- The sketch shows the place from directly above: structures (roof outline), outdoor
  areas, approach paths, and every OPENING (door, gate, window, stair/level access).
- NEVER expose interior floor plans or interior furniture. Indoor action may only be
  represented by its openings (the window/door through which it is visible or
  reached) and by elements visible from outside.
- Every shot must be covered: for each shot, the elements a viewer would need to
  locate that shot's action on the sketch (its opening, area, or approach).
- Keep it compact: at most 10 elements, each with a single capital letter ID
  (A, B, C, ...) and a SHORT English name (1-3 words) for the legend.
- Placement phrases must be relative and physical (e.g. "on the north edge of the
  roof", "at the bottom of the exterior stair"), grounded in the given data only.
- Do not invent features that contradict the member descriptions. Do not use any
  proper names in element names; use generic physical terms."""

R1_SCHEMA = {
    "type": "object",
    "properties": {
        "elements": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "id": {"type": "string", "description": "single capital letter"},
                    "name_en": {"type": "string"},
                    "category": {"type": "string", "enum": [
                        "structure", "opening", "fixture", "approach", "surround"]},
                    "placement": {"type": "string"},
                    "grounding": {"type": "string",
                                  "description": "which shots/member text justify it"},
                },
                "required": ["id", "name_en", "category", "placement", "grounding"],
                "additionalProperties": False,
            },
        },
        "coverage": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "shot_key": {"type": "string"},
                    "element_ids": {"type": "array", "items": {"type": "string"}},
                },
                "required": ["shot_key", "element_ids"],
                "additionalProperties": False,
            },
        },
        "sketch_notes": {"type": "string",
                         "description": "overall massing/orientation guidance for the sketch"},
        "legend_order": {"type": "array", "items": {"type": "string"}},
    },
    "required": ["elements", "coverage", "sketch_notes", "legend_order"],
    "additionalProperties": False,
}


def r1_plan(recon):
    members = F.members_block(recon["members"])
    shot_blocks = "\n\n".join(
        F.shot_block(k, s) for k, s in sorted(recon["shots"].items()))
    user = (
        "PLACE MEMBERS:\n" + members
        + "\n\nALL SHOTS AT THIS PLACE:\n\n" + shot_blocks
        + "\n\nDesign the single aerial site sketch element set now."
    )
    plan = F.llm("forest_root_plan", R1_SYSTEM, user, R1_SCHEMA, model="gpt")
    # 상한 가드(계약 위반 시에도 진행하되 기록)
    if len(plan.get("elements", [])) > 10:
        F.runlog({"kind": "warn", "stage": "s1", "msg":
                  f"elements>{10}: {len(plan['elements'])}"})
    F.save_plan("root_plan", plan)
    covered = {c["shot_key"] for c in plan.get("coverage", [])}
    missing = sorted(set(recon["shots"]) - covered)
    print(f"R1: {len(plan['elements'])} elements, coverage {len(covered)}"
          f"/{len(recon['shots'])} shots" + (f" MISSING={missing}" if missing else ""))
    return plan


def sketch_prompt(plan, members_text: str) -> str:
    el_lines = "\n".join(
        f"({e['id']}) {e['name_en']} — {e['placement']}"
        for e in plan["elements"])
    legend_lines = ", ".join(
        f"{e['id']}: {e['name_en'].upper()}" for e in plan["elements"])
    return "\n".join([
        "ONE hand-drawn TOP-DOWN aerial SITE SKETCH of a single real-world place,",
        "in clean technical sketch style: dark ink outlines, light flat colour",
        "washes, straight-down view. NOT a photograph, NOT a 3D render, no",
        "perspective, no photo textures.",
        "",
        "The place consists of these connected parts — draw them TOGETHER as one",
        "coherent site (one main structure plus its immediate surroundings):",
        members_text,
        "",
        "Overall massing/orientation guidance:",
        plan.get("sketch_notes", ""),
        "",
        "Draw ONLY what is visible from above and from outside: the structure's",
        "roof outline and what stands on top of it, open ground areas, approach",
        "paths and stairs, and each opening (door, gate, window, level access) as",
        "part of the boundary it belongs to. Do NOT draw any interior room layout",
        "or interior furniture — buildings show their roof plane only.",
        "",
        "Mark each of the following elements at its position with a small circle",
        "containing its single capital letter:",
        el_lines,
        "",
        "On the RIGHT side of the image, add a clean legend panel listing each",
        "circled letter with its English name in small neat capital letters:",
        legend_lines,
        "",
        "The circled letters and the legend panel are the ONLY text and the ONLY",
        "circular marks allowed. No other labels, numbers, logos or watermarks.",
        "Include a modest margin of the immediate surroundings so the place reads",
        "in context; keep geometry simple, unambiguous and physically plausible.",
    ])


def r2_images(recon, plan):
    members_text = F.members_block(recon["members"])
    prompt = sketch_prompt(plan, members_text)
    # R-A: 순수 T2I
    F.img_gpt("root_A_t2i", prompt, refs=None, out_path=OUTDIR / "root_A.png")
    # R-B: production 통합 배치도 I2I (기하 유지 + 기호/범례 추가)
    if recon.get("baseline_aerial"):
        edit_prompt = prompt + (
            "\n\nThe attached image is the existing clean top-down site plan of this"
            " place. KEEP its geometry, proportions and layout EXACTLY — same"
            " structures, same openings, same surroundings. Restyle it into the"
            " hand-drawn sketch style described above and ADD the circled capital"
            " letters and the right-side legend panel. Change nothing else."
        )
        F.img_gpt("root_B_from_plan", edit_prompt, refs=[recon["baseline_aerial"]],
                  out_path=OUTDIR / "root_B.png")
    else:
        F.runlog({"kind": "warn", "stage": "s1", "msg": "no baseline aerial for R-B"})


READBACK_SCHEMA = {
    "type": "object",
    "properties": {
        "letters_found": {"type": "array", "items": {"type": "string"}},
        "legend_entries": {"type": "array", "items": {"type": "string"}},
        "interior_layout_exposed": {"type": "boolean"},
        "notes": {"type": "string"},
    },
    "required": ["letters_found", "legend_entries", "interior_layout_exposed",
                 "notes"],
    "additionalProperties": False,
}


def readback(plan, variant: str):
    png = OUTDIR / f"root_{variant}.png"
    if not png.exists():
        return
    user = [
        {"type": "text", "text":
            "This is a top-down site sketch with circled capital letters and a"
            " legend panel. List every circled capital letter you can find, list"
            " the legend entries as written, and say whether any interior room"
            " layout is exposed."},
        F.png_data_url(png),
    ]
    rb = F.llm(f"forest_root_readback", "You read technical sketches precisely.",
               user, READBACK_SCHEMA, model="gpt")
    F.save_plan(f"root_readback_{variant}", rb)
    want = {e["id"] for e in plan["elements"]}
    got = set(rb.get("letters_found") or [])
    print(f"readback {variant}: letters {sorted(got)} vs plan {sorted(want)} | "
          f"interior_exposed={rb.get('interior_layout_exposed')} | {rb.get('notes','')[:120]}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--only", default="all",
                    choices=["all", "plan", "image", "readback"])
    args = ap.parse_args()
    recon = F.load_recon()
    if args.only in ("all", "plan"):
        plan = r1_plan(recon)
    else:
        plan = F.load_plan("root_plan")
    if args.only in ("all", "image"):
        r2_images(recon, plan)
    if args.only in ("all", "readback"):
        for v in ("A", "B"):
            readback(plan, v)
    F.runlog({"kind": "stage", "stage": "s1_root", "done": args.only})


if __name__ == "__main__":
    main()
