"""experiment_floor_plan_image_smoke_slice — W17B.

Single fp_l05_01 smoke: consume the W17A floor-plan image prompt run and
issue ONE gpt-image-2 ``images.generate`` call. The generated PNG is the
intermediate readable diagram that W17C will hand to a VLM for marker
readback (no W17C work in this stage).

Minimal stage by design:
- single fp only (wave-1 lock = fp_l05_01)
- one image, run-local PNG, retry 0
- no DB write, no ImageAsset write, no production manifest mutation
- no commit, no push
- the openai caller is the monkeypatch seam for tests; the script never
  silently retries.

CLI:
  --derive-image-smoke-from <W17A_run_dir>      (required)
  --target-fp-ids fp_l05_01                     (default, restricted)
  --generate                                     (default off → dry-run)
  --model gpt-image-2                            (default; fail-closed)
  --output-root <path>
  --diag-print-imports
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Set

_REPO_ROOT = Path(__file__).resolve().parents[2]
_SCRIPTS_DIR = _REPO_ROOT / "backend" / "scripts"
if str(_SCRIPTS_DIR) not in sys.path:
    sys.path.insert(0, str(_SCRIPTS_DIR))

from experiment_background_pipeline_slice import (  # type: ignore
    KST,
    PLAN_VERSION,
    _check_image_imports_present,
    _check_production_diff_empty,
    _load_backend_env,
    _maybe_print_imports,
)
from experiment_floor_plan_grid_layout_slice import (  # type: ignore
    W16_ALLOWED_TARGET_FP_IDS,
    _resolve_targets,
)
from experiment_actual_floor_plan_generation_slice import (  # type: ignore
    _create_openai_client,
    _generate_png_via_openai,
)

W17B_STAGE = "w17b_floor_plan_image_smoke_slice"
W17B_IMAGE_BACKEND = "gpt-image-2"
W17B_IMAGE_SIZE = "1024x1024"
W17B_IMAGE_QUALITY = "high"

# W17B2 — deterministic marker mapping prepended to the W17A prompt so
# the image model cannot invent its own legend.
# W17B3 extension: each mapping line also carries `unit=<unit_id_pointer>`
# and `placement=<position_hint>` joined from the W15e candidate (chained
# via W17A run_meta.args.derive_image_prompt_from), and the banner adds
# room-local placement guidance so evidence/surface-cue markers stay
# inside the unit the candidate assigned them to.
W17B2_MAPPING_BANNER = (
    "Marker mapping contract (HIGHEST PRIORITY): use exactly these "
    "marker meanings. Do not create a new legend, do not rename "
    "marker labels, do not reuse one marker number for multiple "
    "elements. Prefer NO independent side legend; place each marker "
    "number directly next to its referenced element on the plan. If a "
    "side legend strip is unavoidable, it must copy this mapping "
    "verbatim. Place each marker inside or on the boundary of its "
    "stated unit; prefer room-local placement over global aesthetics. "
    "Do not move evidence or surface-cue markers into another unit. "
    "The plan-marker mapping below is authoritative:"
)


def _load_w15e_candidate_index(
    *, w17a_run_dir: Path, w17a_run_meta: Dict[str, Any],
) -> Dict[str, Dict[int, Dict[str, str]]]:
    """Read the W15e `floor_plan_prompt_candidate.json` referenced by the
    W17A run's `args.derive_image_prompt_from` (or fallback to
    `derived_from` resolved against the topology-slice parent dir) and
    return a `{fp_id: {number: {unit, placement}}}` index. Missing
    candidate file yields an empty index — the script still proceeds
    (mapping lines simply omit `unit=` / `placement=` columns)."""
    candidate_path: Optional[Path] = None
    args_dict = (w17a_run_meta or {}).get("args") or {}
    raw = args_dict.get("derive_image_prompt_from") if isinstance(args_dict, dict) else None
    if isinstance(raw, str) and raw:
        p = Path(raw)
        if not p.is_absolute():
            p = _REPO_ROOT / raw
        if (p / "floor_plan_prompt_candidate.json").exists():
            candidate_path = p / "floor_plan_prompt_candidate.json"
    if candidate_path is None:
        derived_from = (w17a_run_meta or {}).get("derived_from") or ""
        if isinstance(derived_from, str) and derived_from:
            parent = (
                _REPO_ROOT / "scripts_output" / "floor_plan_topology_slice_experiment"
            )
            if (parent / derived_from / "floor_plan_prompt_candidate.json").exists():
                candidate_path = parent / derived_from / "floor_plan_prompt_candidate.json"
    out: Dict[str, Dict[int, Dict[str, str]]] = {}
    if candidate_path is None or not candidate_path.exists():
        return out
    try:
        candidate = json.loads(candidate_path.read_text())
    except Exception:  # noqa: BLE001
        return out
    for fp_id, fp in (candidate.get("candidate_floor_plans") or {}).items():
        if not isinstance(fp, dict):
            continue
        per_num: Dict[int, Dict[str, str]] = {}
        for entry in fp.get("candidate_numbered_elements") or []:
            if not isinstance(entry, dict):
                continue
            n = entry.get("number")
            if not isinstance(n, int):
                continue
            per_num[n] = {
                "unit": entry.get("unit_id_pointer") or "",
                "placement": entry.get("position_hint") or "",
            }
        out[fp_id] = per_num
    return out


def _build_assembled_image_prompt(*, base_prompt: str,
                                  legend: List[Dict[str, Any]],
                                  placement_by_number: Optional[Dict[int, Dict[str, str]]] = None,
                                  ) -> str:
    """Prepend a deterministic `#N = label | kind=... | visual=...
    [| unit=... | placement=...]` mapping block to the W17A
    `t2i_prompt_text`. unit/placement columns are added when the
    W15e candidate carries values for that marker number (joined by
    exact integer match — no semantic judgement)."""
    lines: List[str] = [W17B2_MAPPING_BANNER]
    sortable: List[Dict[str, Any]] = []
    for it in legend or []:
        if isinstance(it, dict) and isinstance(it.get("marker_number"), int):
            sortable.append(it)
    p_by_n = placement_by_number or {}
    for it in sorted(sortable, key=lambda x: x["marker_number"]):
        n = it["marker_number"]
        label = it.get("label") or ""
        kind = it.get("element_kind") or ""
        visual = it.get("visual_encoding") or ""
        cols: List[str] = [
            f"#{n} = {label}", f"kind={kind}", f"visual={visual}",
        ]
        per = p_by_n.get(n) or {}
        unit = per.get("unit") or ""
        placement = per.get("placement") or ""
        if unit:
            cols.append(f"unit={unit}")
        if placement:
            cols.append(f"placement={placement}")
        lines.append(" | ".join(cols))
    block = "\n".join(lines)
    base = (base_prompt or "").rstrip()
    if not base:
        return block
    return block + "\n\n" + base

_DEFAULT_OUTPUT_ROOT = (
    _REPO_ROOT / "scripts_output" / "floor_plan_image_smoke_slice_experiment"
)

# Module-level caller alias is the test monkeypatch seam. Tests replace
# `_openai_caller` with a fake; production calls go through the imported
# helper directly.
_openai_caller: Callable = _generate_png_via_openai


def _run_id() -> str:
    import secrets

    return datetime.now(KST).strftime("%Y%m%d_%H%M") + "_" + secrets.token_hex(3)


def _parse_args(argv):
    p = argparse.ArgumentParser(
        description=(
            "W17B floor-plan image smoke slice — single fp, single image, "
            "gpt-image-2. Default off; --generate to issue the API call."
        )
    )
    p.add_argument(
        "--derive-image-smoke-from", required=True,
        help="Path to a prior W17A success run dir (containing "
             "floor_plan_image_prompt.json + run_meta.json).",
    )
    p.add_argument(
        "--target-fp-ids", default="fp_l05_01",
        help="Comma-separated fp_id subset. Wave 1 only allows fp_l05_01.",
    )
    p.add_argument(
        "--generate", action="store_true",
        help="Actually call gpt-image-2 once. Default off → dry-run.",
    )
    p.add_argument("--model", default=W17B_IMAGE_BACKEND)
    p.add_argument("--output-root", default=str(_DEFAULT_OUTPUT_ROOT))
    p.add_argument("--diag-print-imports", action="store_true")
    return p.parse_args(argv)


def _load_w17a_artifacts(prev_run_dir: Path) -> dict:
    out: Dict[str, Any] = {"_missing": []}
    for key, fname in (
        ("prompt", "floor_plan_image_prompt.json"),
        ("run_meta", "run_meta.json"),
    ):
        p = prev_run_dir / fname
        if not p.exists():
            out["_missing"].append(fname)
            continue
        out[key] = json.loads(p.read_text())
    out["_prev_run_id"] = prev_run_dir.name
    return out


def main(argv=None) -> int:
    args = _parse_args(argv)
    if args.generate:
        _load_backend_env()
    run_id = _run_id()
    out_root = Path(args.output_root)
    run_dir = out_root / run_id
    run_dir.mkdir(parents=True, exist_ok=True)

    prev_run_dir = Path(args.derive_image_smoke_from)
    if not prev_run_dir.is_absolute():
        prev_run_dir = Path.cwd() / prev_run_dir

    artifacts = _load_w17a_artifacts(prev_run_dir)
    missing = list(artifacts.get("_missing", []))

    target_fp_ids, invalid_targets = _resolve_targets(args.target_fp_ids)

    failed_invariants: List[str] = []
    run_status = "succeeded"
    exit_code = 0
    stage_status = "dry_run"
    api_call_count = 0
    image_generation_count = 0

    run_meta: Dict[str, Any] = {
        "run_id": run_id,
        "stage": W17B_STAGE,
        "plan_version": PLAN_VERSION,
        "generated_at": datetime.now(KST).isoformat(),
        "args": vars(args),
        "derived_from": prev_run_dir.name,
        "target_fp_ids": sorted(target_fp_ids),
        "invalid_targets": invalid_targets,
        "model": args.model,
        "image_api_call_count": api_call_count,
        "image_generation_count": image_generation_count,
        "outputs": [],
        "run_status": run_status,
        "exit_code": exit_code,
        "failed_invariants": failed_invariants,
        "stage_status": stage_status,
    }

    def _persist_run_meta_and_exit(code: int) -> int:
        run_meta["run_status"] = (
            "validation_failed" if code != 0 else run_meta["run_status"]
        )
        run_meta["exit_code"] = code
        run_meta["failed_invariants"] = failed_invariants
        (run_dir / "run_meta.json").write_text(
            json.dumps(run_meta, ensure_ascii=False, indent=2)
        )
        return code

    if missing:
        failed_invariants.append("w17a_inputs_missing")
        run_meta["missing_inputs"] = missing
        return _persist_run_meta_and_exit(1)
    if invalid_targets:
        failed_invariants.append("invalid_target_fp_ids")
        run_meta["invalid_targets"] = invalid_targets
        return _persist_run_meta_and_exit(1)
    if not target_fp_ids.issubset(W16_ALLOWED_TARGET_FP_IDS):
        failed_invariants.append("target_fp_outside_allowed")
        return _persist_run_meta_and_exit(1)

    fp_id = sorted(target_fp_ids)[0]
    by_fp = (artifacts["prompt"] or {}).get("floor_plan_image_prompt_by_fp") or {}
    fp_entry = by_fp.get(fp_id) or {}
    prompt_text = fp_entry.get("t2i_prompt_text") or ""
    legend = fp_entry.get("numbered_marker_legend") or []
    if not prompt_text:
        failed_invariants.append("w17a_prompt_text_missing")
        return _persist_run_meta_and_exit(1)

    w15e_index = _load_w15e_candidate_index(
        w17a_run_dir=prev_run_dir, w17a_run_meta=artifacts.get("run_meta") or {},
    )
    placement_by_number = w15e_index.get(fp_id) or {}
    assembled_prompt = _build_assembled_image_prompt(
        base_prompt=prompt_text, legend=legend,
        placement_by_number=placement_by_number,
    )

    png_path = run_dir / "png" / f"{fp_id}.png"
    result: Dict[str, Any] = {
        "fp_id": fp_id,
        "model": args.model,
        "size": W17B_IMAGE_SIZE,
        "quality": W17B_IMAGE_QUALITY,
        "source_prompt_len_chars": len(prompt_text),
        "assembled_prompt_len_chars": len(assembled_prompt),
        "assembled_prompt": assembled_prompt,
        "legend_entry_count": len(legend),
        "placement_carry_count": len(placement_by_number),
        "png_relative_path": str(png_path.relative_to(run_dir)),
        "status": "dry_run",
    }

    if args.generate:
        if (args.model or "").strip() != W17B_IMAGE_BACKEND:
            failed_invariants.append("model_must_be_gpt_image_2")
            run_meta["model_error"] = f"got {args.model!r}, required {W17B_IMAGE_BACKEND!r}"
            return _persist_run_meta_and_exit(1)
        if not os.environ.get("OPENAI_API_KEY"):
            failed_invariants.append("openai_api_key_missing")
            return _persist_run_meta_and_exit(1)
        try:
            client = _create_openai_client()
        except Exception as exc:  # noqa: BLE001
            failed_invariants.append("openai_client_unavailable")
            run_meta["openai_client_error"] = str(exc)[:240]
            return _persist_run_meta_and_exit(1)

        api_call_count = 1
        outcome = _openai_caller(
            fp_id=fp_id, prompt=assembled_prompt, model=args.model,
            size=W17B_IMAGE_SIZE, quality=W17B_IMAGE_QUALITY,
            target_path=png_path, client=client,
        )
        result.update(outcome)
        if outcome.get("status") == "success":
            stage_status = "generated"
            image_generation_count = 1
        else:
            failed_invariants.append("openai_call_failed")
            run_status = "validation_failed"
            exit_code = 1
            stage_status = "api_call_failed"

    (run_dir / "generated_image_meta.json").write_text(
        json.dumps({"fp_smoke_result": result}, ensure_ascii=False, indent=2)
    )
    run_meta["outputs"].append("generated_image_meta.json")

    # Light safety surface (kept lean per Codex: minimal new tests, light
    # report). The single-run JSON file documents the api call shape.
    safety: Dict[str, Any] = {
        "production_diff_empty": _check_production_diff_empty(),
        "image_imports_seen": _check_image_imports_present(),
        "api_call_count": api_call_count,
        "image_generation_count": image_generation_count,
        "image_api_call_at_most_one":  api_call_count <= 1,
        "target_fp_only_allowed": target_fp_ids.issubset(W16_ALLOWED_TARGET_FP_IDS),
        "png_emitted_when_generated": (
            png_path.exists() and stage_status == "generated"
        ) or stage_status != "generated",
    }
    (run_dir / "safety_surface.json").write_text(
        json.dumps(safety, ensure_ascii=False, indent=2)
    )
    run_meta["outputs"].append("safety_surface.json")

    if not safety["production_diff_empty"]:
        failed_invariants.append("production_diff_dirty")
        run_status = "validation_failed"
        exit_code = 1
    if not safety["image_api_call_at_most_one"]:
        failed_invariants.append("image_api_call_count_exceeds_one")
        run_status = "validation_failed"
        exit_code = 1

    run_meta["image_api_call_count"] = api_call_count
    run_meta["image_generation_count"] = image_generation_count
    run_meta["stage_status"] = stage_status
    run_meta["run_status"] = run_status
    run_meta["exit_code"] = exit_code
    run_meta["failed_invariants"] = failed_invariants
    run_meta["safety_surface"] = safety
    run_meta["png_relative_path"] = result["png_relative_path"]
    run_meta["png_size_bytes"] = (
        png_path.stat().st_size if png_path.exists() else 0
    )

    # Light HTML preview for the smoke run.
    def esc(x):
        return (
            str(x).replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
        )

    png_html = (
        f'<img src="{esc(result["png_relative_path"])}" '
        f'style="max-width:720px;border:1px solid #ccc"/>'
        if png_path.exists() else "<p><i>(no PNG — dry-run or failed call)</i></p>"
    )
    (run_dir / "index.html").write_text(
        f"""<!doctype html><html><head><meta charset=\"utf-8\">
<title>W17B floor_plan_image_smoke_slice {esc(run_id)}</title>
<style>body{{font-family:sans-serif;margin:1.5em}}
pre{{white-space:pre-wrap;font-size:0.85em;max-width:90ch}}</style></head>
<body>
<h1>W17B — floor_plan_image_smoke_slice {esc(run_id)}</h1>
<p>stage: <b>{esc(stage_status)}</b>
| run_status: <b>{esc(run_status)}</b>
| exit_code: {esc(exit_code)}
| model: <b>{esc(args.model)}</b>
| api_call_count: <b>{esc(api_call_count)}</b>
| image_generation_count: <b>{esc(image_generation_count)}</b>
| png_size_bytes: <b>{esc(run_meta['png_size_bytes'])}</b></p>
{png_html}
<details><summary>generated_image_meta.json</summary>
<pre>{esc(json.dumps(result, ensure_ascii=False, indent=2))}</pre></details>
<details><summary>safety_surface.json</summary>
<pre>{esc(json.dumps(safety, ensure_ascii=False, indent=2))}</pre></details>
<details><summary>run_meta.json</summary>
<pre>{esc(json.dumps(run_meta, ensure_ascii=False, indent=2))}</pre></details>
</body></html>"""
    )
    run_meta["outputs"].append("index.html")
    if png_path.exists():
        run_meta["outputs"].append(str(png_path.relative_to(run_dir)))

    (run_dir / "run_meta.json").write_text(
        json.dumps(run_meta, ensure_ascii=False, indent=2)
    )
    _maybe_print_imports(args)
    return exit_code


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