"""s2 — 실내 체인: I0 계획 → I1 탑레이어 → I2 fp변형 → 임시샷 → A/B/C 경로 → 최종.

사용: .venv/bin/python s2_indoor.py [--only plan|toplayer|fpvar|temp|paths]
                                    [--alpha-shot 첫샷 임시배경 α 변형] [--mannequin-preview]
"""
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
import forest_chain as C  # noqa: E402

OUTDIR = F.OUT / "indoor"

I0_SYSTEM = """You are the background production planner for the INDOOR shots of one
real-world place. You are given: the place's aerial-sketch element plan (the spatial
root contract), the indoor floor plan identifiers available, and the FULL content of
every indoor shot.

Plan the background generation like a set photographer:
1. For each shot, state what background it needs (which part of the indoor space,
   from roughly what direction).
2. Merge those needs into AT MOST TWO "top-layer" master backgrounds — wide, clean,
   empty-room views that together cover all the shots. Prefer ONE if the space is
   small enough to read in a single wide view.
3. Give the generation order. If there are two, the second must chain on the first
   (same dwelling, same materials).
4. For each top-layer, judge whether attaching the floor plan image as a reference
   would help (geometry anchoring) or hurt (style leakage), and say why.
Ground every statement in the given data only; do not invent rooms or furniture.
Use generic physical terms, no proper names."""

I0_SCHEMA = {
    "type": "object",
    "properties": {
        "top_layers": {"type": "array", "items": {"type": "object", "properties": {
            "key": {"type": "string"},
            "loc_id": {"type": "string"},
            "purpose": {"type": "string"},
            "view_description": {"type": "string",
                                 "description": "what the master background shows, "
                                                "phrased as an empty-room camera view"},
            "attach_fp": {"type": "boolean"},
            "attach_fp_reason": {"type": "string"},
            "order": {"type": "integer"},
            "chain_prev_key": {"type": ["string", "null"]}},
            "required": ["key", "loc_id", "purpose", "view_description",
                         "attach_fp", "attach_fp_reason", "order",
                         "chain_prev_key"],
            "additionalProperties": False}},
        "shot_assignments": {"type": "array", "items": {"type": "object",
            "properties": {
                "shot_key": {"type": "string"},
                "top_layer_key": {"type": "string"},
                "needed_background": {"type": "string"}},
            "required": ["shot_key", "top_layer_key", "needed_background"],
            "additionalProperties": False}},
    },
    "required": ["top_layers", "shot_assignments"],
    "additionalProperties": False,
}


def i0_plan(recon, root_plan):
    indoor_shots = {k: v for k, v in recon["shots"].items() if v["is_indoor"]}
    fp_lines = "\n".join(
        f"- loc {loc}: floor plan(s) {', '.join(x['fp_id'] for x in fps)}"
        for loc, fps in recon["fp_by_loc"].items())
    user = (
        "PLACE MEMBERS:\n" + F.members_block(recon["members"])
        + "\n\nAERIAL ROOT CONTRACT (element plan):\n"
        + json.dumps(root_plan, ensure_ascii=False)
        + "\n\nAVAILABLE FLOOR PLANS:\n" + fp_lines
        + "\n\nALL INDOOR SHOTS:\n\n"
        + "\n\n".join(F.shot_block(k, v) for k, v in sorted(indoor_shots.items()))
        + "\n\nPlan the top-layer backgrounds now."
    )
    plan = F.llm("forest_indoor_plan", I0_SYSTEM, user, I0_SCHEMA, model="gpt")
    F.save_plan("indoor_plan", plan)
    print(f"I0: {len(plan['top_layers'])} top layers; "
          f"{len(plan['shot_assignments'])} shot assignments")
    for tl in plan["top_layers"]:
        print(f"  [{tl['order']}] {tl['key']} loc={tl['loc_id']} "
              f"attach_fp={tl['attach_fp']} chain={tl['chain_prev_key']}")
    return plan


def top_layer_prompt(recon, tl, root_plan, chained: bool) -> str:
    loc = next((m for m in recon["members"] if m["loc_id"] == tl["loc_id"]), {})
    lines = [
        "Create ONE photorealistic EMPTY interior master background plate of a",
        "real dwelling interior — a wide, clean view with NO people, no text, no",
        "markers. Real-world human scale (doorways adult height).",
        "",
        f"THE PLACE (from production data): {loc.get('label')}: {loc.get('summary')}",
        f"THIS PLATE'S PURPOSE: {tl['purpose']}",
        f"THE VIEW: {tl['view_description']}",
        "",
        "Site orientation notes from the aerial root contract:",
        str(root_plan.get("sketch_notes", "")),
    ]
    if chained:
        lines += [
            "",
            "The LAST attached image shows ANOTHER part of THE SAME dwelling —",
            "match its wall colours, floor material, trim, door style, lighting",
            "mood and renovation age EXACTLY so both read as the same home. Do",
            "NOT copy its camera angle or composition.",
        ]
    return "\n".join(lines)


FP_CLAUSE = (
    "\n\nThe FIRST attached image is the TOP-DOWN floor plan of this dwelling —"
    " use it ONLY to get the room geometry, openings and furniture positions"
    " right. NEVER copy its drawing style, numbered circles or any marks into"
    " the photorealistic output."
)


def i1_top_layers(recon, plan, root_plan, force_ab: bool):
    prev_png = None
    ordered = sorted(plan["top_layers"], key=lambda t: t["order"])
    for idx, tl in enumerate(ordered):
        fp_list = recon["fp_by_loc"].get(tl["loc_id"]) or []
        fp_png = fp_list[0]["png_path"] if fp_list else None
        chained = bool(tl.get("chain_prev_key")) and prev_png is not None
        prompt = top_layer_prompt(recon, tl, root_plan, chained)
        refs = []
        use_fp = bool(tl.get("attach_fp")) and fp_png
        if use_fp:
            prompt_full = prompt + FP_CLAUSE
            refs.append(fp_png)
        else:
            prompt_full = prompt
        if chained:
            refs.append(prev_png)
        out = OUTDIR / "top_layers" / f"{tl['key']}.png"
        F.img_gpt(f"top_{tl['key']}", prompt_full, refs=refs or None,
                  size="1536x1024", out_path=out)
        prev_png = out

        # 강제 fp on/off A/B — 첫 탑레이어 1쌍만 (Codex 의견 수용: 원인 분리)
        if force_ab and idx == 0 and fp_png:
            if not use_fp:
                F.img_gpt(f"top_{tl['key']}_fpON", prompt + FP_CLAUSE,
                          refs=[fp_png], size="1536x1024",
                          out_path=OUTDIR / "top_layers" / f"{tl['key']}_fpON.png")
            else:
                F.img_gpt(f"top_{tl['key']}_fpOFF", prompt, refs=None,
                          size="1536x1024",
                          out_path=OUTDIR / "top_layers" / f"{tl['key']}_fpOFF.png")


BLOCKING_PROMPT = (
    "The attached image is a TOP-DOWN floor plan of one dwelling: numbered circles"
    " mark fixed rooms and furniture. Keep the plan EXACTLY as it is — same walls,"
    " rooms, furniture and numbered circles, do not redraw or restyle anything.\n"
    "ADD exactly these annotations for ONE camera setup:\n"
    "- ONE green camera icon with a green triangular view-CONE showing where the"
    " camera stands and what it looks at, following this staging description:\n"
    "{cam}\n"
    "{figures}"
    "Add NOTHING else — no new rooms, no new furniture, no text other than the"
    " existing numbers and the new capital letters."
)


def figures_clause(s) -> str:
    st = s.get("staging") or {}
    cons = ((st.get("frame_spatial_contract") or {}).get("constraints") or [])
    n = sum(1 for c in cons if isinstance(c, dict)
            and c.get("target_kind") == "character")
    if n <= 0:
        n = len(st.get("character_angles") or [])
    if n <= 0:
        return "- This shot has NO people — add no lettered circles.\n"
    letters = " ".join(f"({chr(65 + i)})" for i in range(n))
    return (f"- {n} red circle(s) with capital letters {letters} inside the"
            f" view-cone where the staging places the people:\n"
            + "\n".join("  * " + json.dumps(ca, ensure_ascii=False)
                        for ca in (st.get("character_angles") or [])) + "\n")


def i2_fpvar(recon, sk):
    s = recon["shots"][sk]
    loc = s["loc_ids"][0]
    fp_list = recon["fp_by_loc"].get(loc) or []
    if not fp_list:
        raise SystemExit(f"no fp for indoor shot {sk} loc {loc}")
    cam = str((s.get("staging") or {}).get("camera_direction") or "")
    prompt = BLOCKING_PROMPT.format(cam=cam, figures=figures_clause(s))
    out = OUTDIR / sk / "fpvar.png"
    return F.img_gpt(f"{sk}_fpvar", prompt, refs=[fp_list[0]["png_path"]],
                     size="1024x1024", out_path=out)


def assigned_top(plan, sk) -> str:
    for a in plan["shot_assignments"]:
        if a["shot_key"] == sk:
            return a["top_layer_key"]
    return sorted(plan["top_layers"], key=lambda t: t["order"])[0]["key"]


def run_shot(recon, plan, sk, alpha: bool, mannequin_preview: bool):
    s = recon["shots"][sk]
    d = OUTDIR / sk
    top_key = assigned_top(plan, sk)
    env = OUTDIR / "top_layers" / f"{top_key}.png"
    fpvar = d / "fpvar.png"
    passports = C.passports_for_shot(recon, s)

    # 임시샷 (β: 탑레이어 직참조 — 기본)
    temp = C.temp_shot(f"{sk}_temp", s, env, fpvar, d / "temp.png")

    if alpha:  # α: 임시배경(샷 맞춤 빈 배경) 경유 — 실험 변수
        tb_prompt = "\n".join([
            "Create the photorealistic EMPTY interior background seen by the",
            "camera marked on the attached top-down diagram (green icon+cone),",
            "at eye level, real-world scale. Use the other attached image as the",
            "material/lighting ground truth of the same dwelling. NO people, no",
            "text, no markers.",
            "", C.shot_content_lines(s), C.staging_lines(s),
        ])
        tb = F.img_gpt(f"{sk}_tempbg", tb_prompt, refs=[fpvar, env],
                       out_path=d / "temp_bg.png")
        C.temp_shot(f"{sk}_temp_alpha", s, tb, fpvar, d / "temp_alpha.png")

    # 경로 A: 직행
    C.final_shot(f"{sk}_final_A", s, env, passports, d / "final_A.png",
                 control_png=fpvar)
    # 경로 B: VLM 구도 텍스트 → 배경 재생성 → 최종
    comp = C.vlm_composition(sk, temp)
    bg_b = C.bg_regen(f"{sk}_bgB", s, env, d / "bg_B.png", comp=comp)
    C.final_shot(f"{sk}_final_B", s, bg_b, passports, d / "final_B.png")
    # 경로 C: i2i 윤곽 → 배경 재생성 → 최종 (사용자 지정 주경로)
    cont = C.contour(f"{sk}_contour", temp, d / "contour.png")
    bg_c = C.bg_regen(f"{sk}_bgC", s, env, d / "bg_C.png", contour_png=cont)
    C.final_shot(f"{sk}_final_C", s, bg_c, passports, d / "final_C.png")

    if mannequin_preview:
        bg_cm = C.bg_regen(f"{sk}_bgC_m", s, env, d / "bg_C_m.png",
                           contour_png=cont, mannequin=True)
        C.final_shot(f"{sk}_final_C_m", s, bg_cm, passports, d / "final_C_m.png")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--only", default="all",
                    choices=["all", "plan", "toplayer", "fpvar", "temp", "paths"])
    ap.add_argument("--no-force-ab", action="store_true")
    args = ap.parse_args()
    recon = F.load_recon()
    root_plan = F.load_plan("root_plan")
    subset = F.load_plan("subset")["indoor"]

    if args.only in ("all", "plan"):
        plan = i0_plan(recon, root_plan)
    else:
        plan = F.load_plan("indoor_plan")
    if args.only in ("all", "toplayer"):
        i1_top_layers(recon, plan, root_plan, force_ab=not args.no_force_ab)
    if args.only in ("all", "fpvar", "temp", "paths"):
        for i, sk in enumerate(subset):
            i2_fpvar(recon, sk)
            if args.only in ("all", "temp", "paths"):
                run_shot(recon, plan, sk, alpha=(i == 0),
                         mannequin_preview=(i == 0))
    F.runlog({"kind": "stage", "stage": "s2_indoor", "done": args.only,
              "subset": subset})


if __name__ == "__main__":
    main()
