"""W18I consistency-reference background smoke.

This experimental slice fixes the W18G/H failure mode where every BG was
generated independently from only the base floor-plan reference. W18I plans a
small reference graph:

- first generate a broad clean home-style anchor (default L05B02) from the base
  FP plus the W18F prompt text,
- then generate related BGs from at most two prior visual BG references, not
  from the FP, so material/style identity can carry across views.

Default is dry-run. No DB write, no ImageAsset write, no production mutation.
"""
from __future__ import annotations

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

_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_actual_floor_plan_generation_slice import (  # type: ignore
    _create_openai_client,
    _to_jsonable,
)
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_background_image_smoke_from_region_ref_slice import (  # type: ignore
    W18G_IMAGE_BACKEND,
    W18G_IMAGE_SIZE,
    _resolve_w18b_base_fp_png,
)


W18I_STAGE = "w18i_background_image_consistency_reference_slice"
W18I_ALLOWED_BG_IDS = {"L05B01", "L05B02", "L05B03", "L05B04", "L05B05"}
W18I_GRAPH_ORDER = ["L05B02", "L05B05", "L05B04", "L05B03", "L05B01"]

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

_MODEST_STYLE_PREAMBLE = (
    "Create a compact, modest, ordinary domestic interior. Keep the home "
    "identity practical and lived-in, not luxurious: no showroom, no hotel, "
    "no marble, no chandelier, no designer furniture, no glossy real-estate "
    "brochure styling. Use simple off-white walls, modest light wood or worn "
    "vinyl floor, plain fabric window coverings, ordinary small fixtures, and a "
    "single consistent warm indoor light. Preserve the region and markers "
    "from the prompt. Do not render any person, body, or corpse."
)

_REFERENCE_GUIDANCE_BY_BG = {
    "L05B02": (
        "This is the broad clean home-style anchor. Use the base FP reference "
        "only for layout guidance; establish the ordinary material palette "
        "and fixed furniture identity for later related views."
    ),
    "L05B05": (
        "Inherit the home identity from the broad clean anchor. "
        "Create the clean private-room view with matching floor, wall, window-covering, "
        "and lighting palette."
    ),
    "L05B04": (
        "Keep the same private-room identity as the clean private-room anchor. "
        "Only add the transient cues requested by this BG prompt."
    ),
    "L05B03": (
        "Bridge the broad clean anchor and the private-room anchor. Keep "
        "both visible spaces consistent; only add the transient cues requested "
        "by this BG prompt."
    ),
    "L05B01": (
        "Inherit the broad home palette from the broad clean anchor. "
        "Create the adjacent service-room view with matching ordinary domestic "
        "finishes; only add the requested transient cue."
    ),
}


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=(
            "W18I consistency-reference background smoke. Dry-run by default; "
            "when --generate, issue one images.edit call per target in graph "
            "order. Default target is L05B02 anchor only."
        )
    )
    p.add_argument(
        "--derive-consistency-from",
        required=True,
        help="Path to a prior W18F success run dir.",
    )
    p.add_argument(
        "--target-bg-ids",
        default="L05B02",
        help=(
            "CSV bg_ids. Default is the first clean broad anchor only. "
            "Allowed: L05B01,L05B02,L05B03,L05B04,L05B05."
        ),
    )
    p.add_argument("--generate", action="store_true")
    p.add_argument("--model", default=W18G_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 _resolve_target_bg_ids(raw: str) -> tuple[list[str], list[str]]:
    parts = [p.strip() for p in (raw or "").split(",") if p.strip()]
    valid: List[str] = []
    invalid: List[str] = []
    for bg_id in W18I_GRAPH_ORDER:
        if bg_id in parts and bg_id not in valid:
            valid.append(bg_id)
    for p in parts:
        if p not in W18I_ALLOWED_BG_IDS:
            invalid.append(p)
    return valid, invalid


def _repo_path(path_text: str) -> Path:
    p = Path(path_text)
    if not p.is_absolute():
        p = _REPO_ROOT / p
    return p


def _planned_ref_specs(
    *, bg_id: str, run_dir: Path, base_fp_png_path: Optional[Path]
) -> List[Dict[str, Any]]:
    """Return reference specs in the order sent to images.edit.

    Only the broad first anchor uses the FP reference. Subsequent BGs use up
    to two previously generated visual BG references.
    """
    if bg_id == "L05B02":
        return [{
            "kind": "base_fp_layout",
            "role": "layout_only_for_broad_anchor",
            "path": str(base_fp_png_path) if base_fp_png_path else "",
        }]
    if bg_id == "L05B05":
        return [{
            "kind": "visual_anchor",
            "role": "primary_home_style_anchor",
            "source_bg_id": "L05B02",
            "path": str(run_dir / "png" / "L05B02.png"),
        }]
    if bg_id == "L05B04":
        return [{
            "kind": "visual_anchor",
            "role": "same_private_room_clean_anchor",
            "source_bg_id": "L05B05",
            "path": str(run_dir / "png" / "L05B05.png"),
        }]
    if bg_id == "L05B03":
        return [
            {
                "kind": "visual_anchor",
                "role": "broad_home_style_anchor",
                "source_bg_id": "L05B02",
                "path": str(run_dir / "png" / "L05B02.png"),
            },
            {
                "kind": "visual_anchor",
                "role": "private_room_anchor",
                "source_bg_id": "L05B05",
                "path": str(run_dir / "png" / "L05B05.png"),
            },
        ]
    if bg_id == "L05B01":
        return [{
            "kind": "visual_anchor",
            "role": "primary_home_style_anchor",
            "source_bg_id": "L05B02",
            "path": str(run_dir / "png" / "L05B02.png"),
        }]
    return []


def _build_consistency_prompt(*, bg_id: str, assembled_preview: str) -> str:
    guidance = _REFERENCE_GUIDANCE_BY_BG.get(bg_id, "")
    return (
        _MODEST_STYLE_PREAMBLE
        + "\n\nReference strategy for this BG: "
        + guidance
        + "\n\n"
        + assembled_preview
    )


def _edit_png_via_openai_with_references(
    *, bg_id: str, prompt: str, model: str, size: str,
    reference_paths: List[Path], target_path: Path, client: Any,
) -> Dict[str, Any]:
    t0 = time.monotonic()
    handles = []
    try:
        for ref in reference_paths:
            handles.append(open(ref, "rb"))
        resp = client.images.edit(
            model=model, image=handles, prompt=prompt, size=size, n=1,
        )
    except Exception as exc:  # noqa: BLE001
        latency_ms = int((time.monotonic() - t0) * 1000)
        return {
            "status": "api_call_failed",
            "png_size_bytes": 0,
            "actual_api_response_meta": {"latency_ms": latency_ms},
            "error_meta": {
                "status_code": getattr(exc, "status_code", None),
                "message": str(exc)[:240],
            },
            "cost_meta": {},
        }
    finally:
        for h in handles:
            try:
                h.close()
            except Exception:
                pass

    latency_ms = int((time.monotonic() - t0) * 1000)
    data = getattr(resp, "data", None) or []
    if not data or not getattr(data[0], "b64_json", None):
        return {
            "status": "api_call_failed",
            "png_size_bytes": 0,
            "actual_api_response_meta": {"latency_ms": latency_ms},
            "error_meta": {"message": "response missing b64_json payload"},
            "cost_meta": {},
        }

    target_path.parent.mkdir(parents=True, exist_ok=True)
    target_path.write_bytes(base64.b64decode(data[0].b64_json))

    api_meta: Dict[str, Any] = {"latency_ms": latency_ms}
    for fname in ("created", "request_id", "id", "model"):
        val = getattr(resp, fname, None)
        if val is not None:
            api_meta[fname] = _to_jsonable(val)
    cost_meta: Dict[str, Any] = {}
    usage = getattr(resp, "usage", None)
    if usage is not None:
        cost_meta["provider_usage"] = _to_jsonable(usage)

    return {
        "status": "success",
        "png_size_bytes": target_path.stat().st_size,
        "actual_api_response_meta": api_meta,
        "error_meta": {},
        "cost_meta": cost_meta,
    }


_openai_edit_caller: Callable = _edit_png_via_openai_with_references


def _build_report(
    *, target_bg_ids: List[str], invalid_targets: List[str],
    per_bg_results: Dict[str, Dict[str, Any]], mode: str,
    image_api_call_count: int, image_generation_count: int,
    production_diff_empty: bool, image_import_seen: bool,
) -> Dict[str, Any]:
    inv: Dict[str, Dict[str, Any]] = {}
    inv["inputs_present"] = {
        "pass": bool(per_bg_results),
        "detail": {"per_bg_result_count": len(per_bg_results)},
    }
    inv["target_bgs_supported"] = {
        "pass": bool(target_bg_ids) and not invalid_targets,
        "detail": {
            "target_bg_ids": target_bg_ids,
            "invalid_targets": invalid_targets,
            "allowed": sorted(W18I_ALLOWED_BG_IDS),
        },
    }
    graph_failures: List[Dict[str, Any]] = []
    for bg_id, row in per_bg_results.items():
        refs = row.get("reference_specs") or []
        if len(refs) > 2:
            graph_failures.append({"bg_id": bg_id, "reason": "too_many_refs"})
        if bg_id == "L05B02":
            if not refs or refs[0].get("kind") != "base_fp_layout":
                graph_failures.append({
                    "bg_id": bg_id, "reason": "anchor_must_use_base_fp",
                })
        else:
            if any((r or {}).get("kind") == "base_fp_layout" for r in refs):
                graph_failures.append({
                    "bg_id": bg_id,
                    "reason": "non_anchor_must_not_use_base_fp",
                })
    inv["reference_graph_max_two_and_fp_only_for_anchor"] = {
        "pass": not graph_failures,
        "detail": {"failures": graph_failures},
    }
    if mode == "dry_run":
        gen_ok = image_api_call_count == 0 and image_generation_count == 0
    else:
        gen_ok = (
            image_api_call_count == len(per_bg_results)
            and image_generation_count == len(per_bg_results)
            and all((row or {}).get("png_exists_on_disk")
                    for row in per_bg_results.values())
        )
    inv["image_generation_count_matches_mode"] = {
        "pass": gen_ok,
        "detail": {
            "mode": mode,
            "image_api_call_count": image_api_call_count,
            "image_generation_count": image_generation_count,
        },
    }
    inv["production_diff_zero_db_write_zero_no_imageasset_write"] = {
        "pass": production_diff_empty and not image_import_seen,
        "detail": {
            "production_diff_empty": production_diff_empty,
            "db_write_count": 0,
            "image_asset_write_count": 0,
            "image_import_seen": image_import_seen,
        },
    }
    return {"invariants": inv, "all_pass": all(v["pass"] for v in inv.values())}


def _esc(x: Any) -> str:
    return str(x).replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")


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_consistency_from)
    if not prev_run_dir.is_absolute():
        prev_run_dir = Path.cwd() / prev_run_dir

    target_bg_ids, invalid_targets = _resolve_target_bg_ids(args.target_bg_ids)
    failed_invariants: List[str] = []
    mode = "generated" if args.generate else "dry_run"
    image_api_call_count = 0
    image_generation_count = 0
    run_status = "succeeded"
    exit_code = 0

    w18f_path = prev_run_dir / "background_prompt_region_ref_assembly_preview.json"
    missing_inputs: List[str] = []
    if not w18f_path.exists():
        missing_inputs.append("background_prompt_region_ref_assembly_preview.json")
    if invalid_targets:
        failed_invariants.append("invalid_target_bg_ids")
    if not target_bg_ids:
        failed_invariants.append("no_target_bg_ids")
    if missing_inputs:
        failed_invariants.append("w18f_inputs_missing")

    w18f_payload: Dict[str, Any] = {}
    if w18f_path.exists():
        w18f_payload = json.loads(w18f_path.read_text())

    first_bg = target_bg_ids[0] if target_bg_ids else "L05B02"
    base_fp_png_path = _resolve_w18b_base_fp_png(
        w18f_assembly_payload=w18f_payload, bg_id=first_bg,
    )

    client = None
    if args.generate and not failed_invariants:
        if (args.model or "").strip() != W18G_IMAGE_BACKEND:
            failed_invariants.append("model_must_be_gpt_image_2")
        elif not os.environ.get("OPENAI_API_KEY"):
            failed_invariants.append("openai_api_key_missing")
        else:
            try:
                client = _create_openai_client()
            except Exception as exc:  # noqa: BLE001
                failed_invariants.append("openai_client_unavailable")
                client = {"error": str(exc)[:240]}

    per_bg_results: Dict[str, Dict[str, Any]] = {}
    for bg_id in target_bg_ids:
        bg_payload = (
            (w18f_payload or {}).get("assembly_preview_by_bg") or {}
        ).get(bg_id) or {}
        assembled_preview = bg_payload.get("assembled_background_prompt_preview") or ""
        prompt_text = _build_consistency_prompt(
            bg_id=bg_id, assembled_preview=assembled_preview,
        ) if assembled_preview else ""
        reference_specs = _planned_ref_specs(
            bg_id=bg_id, run_dir=run_dir, base_fp_png_path=base_fp_png_path,
        )
        reference_paths = [
            _repo_path(r.get("path") or "")
            for r in reference_specs
            if r.get("path")
        ]
        missing_refs = [
            str(p) for p in reference_paths if not p.exists()
        ]
        png_path = run_dir / "png" / f"{bg_id}.png"
        row: Dict[str, Any] = {
            "bg_id": bg_id,
            "fp_id": bg_payload.get("fp_id") or "",
            "model": args.model,
            "size": W18G_IMAGE_SIZE,
            "status": "dry_run",
            "clean_background_expected": bool(
                bg_payload.get("clean_background_expected")
            ),
            "transient_overlay_marker_numbers": list(
                bg_payload.get("transient_overlay_marker_numbers") or []
            ),
            "reference_specs": reference_specs,
            "reference_paths": [str(p) for p in reference_paths],
            "missing_reference_paths": missing_refs,
            "prompt_text_len_chars": len(prompt_text),
            "style_preamble": _MODEST_STYLE_PREAMBLE,
            "reference_guidance": _REFERENCE_GUIDANCE_BY_BG.get(bg_id, ""),
            "png_relative_path": str(png_path.relative_to(run_dir)),
            "png_exists_on_disk": False,
        }

        if args.generate and not failed_invariants:
            if not prompt_text:
                row["status"] = "skipped_missing_prompt"
                failed_invariants.append(f"prompt_missing_for_{bg_id}")
            elif missing_refs:
                row["status"] = "skipped_missing_references"
                failed_invariants.append(f"refs_missing_for_{bg_id}")
            else:
                image_api_call_count += 1
                outcome = _openai_edit_caller(
                    bg_id=bg_id, prompt=prompt_text, model=args.model,
                    size=W18G_IMAGE_SIZE, reference_paths=reference_paths,
                    target_path=png_path, client=client,
                )
                row.update(outcome)
                if outcome.get("status") == "success":
                    image_generation_count += 1
                    row["png_exists_on_disk"] = png_path.exists()
                else:
                    failed_invariants.append(f"openai_call_failed_for_{bg_id}")
        per_bg_results[bg_id] = row

    production_diff_empty = _check_production_diff_empty()
    image_import_seen = _check_image_imports_present()
    report = _build_report(
        target_bg_ids=target_bg_ids,
        invalid_targets=invalid_targets,
        per_bg_results=per_bg_results,
        mode=mode,
        image_api_call_count=image_api_call_count,
        image_generation_count=image_generation_count,
        production_diff_empty=production_diff_empty,
        image_import_seen=image_import_seen,
    )
    for name, inv in report["invariants"].items():
        if not inv["pass"] and name not in failed_invariants:
            failed_invariants.append(name)
    if failed_invariants:
        run_status = "validation_failed"
        exit_code = 1

    plan = {
        "stage": W18I_STAGE,
        "source_w18f_run_id": prev_run_dir.name,
        "style_contract": _MODEST_STYLE_PREAMBLE,
        "graph_order": W18I_GRAPH_ORDER,
        "per_bg_results": per_bg_results,
        "next_image_gate": (
            "Start with L05B02 anchor only. Review it before cascading "
            "L05B05/L05B04/L05B03/L05B01."
        ),
    }
    (run_dir / "consistency_reference_plan.json").write_text(
        json.dumps(plan, ensure_ascii=False, indent=2)
    )
    (run_dir / "w18i_compatibility_report.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2)
    )

    run_meta = {
        "run_id": run_id,
        "stage": W18I_STAGE,
        "plan_version": PLAN_VERSION,
        "generated_at": datetime.now(KST).isoformat(),
        "args": vars(args),
        "derived_from": prev_run_dir.name,
        "target_bg_ids": target_bg_ids,
        "invalid_targets": invalid_targets,
        "mode": mode,
        "model": args.model,
        "image_api_call_count": image_api_call_count,
        "image_generation_count": image_generation_count,
        "run_status": run_status,
        "exit_code": exit_code,
        "failed_invariants": failed_invariants,
        "outputs": [
            "consistency_reference_plan.json",
            "w18i_compatibility_report.json",
            "index.html",
        ],
    }

    cards = []
    for bg_id in target_bg_ids:
        row = per_bg_results.get(bg_id) or {}
        refs = "".join(
            f"<li>{_esc((r or {}).get('kind'))} / "
            f"{_esc((r or {}).get('role'))}: "
            f"<code>{_esc((r or {}).get('path'))}</code></li>"
            for r in row.get("reference_specs") or []
        )
        png_html = ""
        png_rel = row.get("png_relative_path") or ""
        if (run_dir / png_rel).exists():
            png_html = (
                f'<img src="{_esc(png_rel)}" '
                f'style="max-width:480px;border:1px solid #ccc"/>'
            )
        else:
            png_html = "<p><i>No PNG in dry-run or failed generate.</i></p>"
        cards.append(
            f"<section><h2>{_esc(bg_id)}</h2>"
            f"<p>status: <b>{_esc(row.get('status'))}</b> | "
            f"clean: <b>{_esc(row.get('clean_background_expected'))}</b> | "
            f"transient: <code>{_esc(row.get('transient_overlay_marker_numbers'))}</code></p>"
            f"<h3>reference images</h3><ol>{refs}</ol>"
            f"<h3>style/reference guidance</h3><pre>{_esc(row.get('reference_guidance'))}</pre>"
            f"{png_html}"
            f"</section>"
        )

    inv_rows = "".join(
        f"<tr><td>{_esc(k)}</td>"
        f"<td class=\"{'pass' if v['pass'] else 'fail'}\">"
        f"{'PASS' if v['pass'] else 'FAIL'}</td>"
        f"<td><pre>{_esc(json.dumps(v.get('detail'), ensure_ascii=False))[:900]}</pre></td></tr>"
        for k, v in (report.get("invariants") or {}).items()
    )
    (run_dir / "index.html").write_text(
        f"""<!doctype html><html><head><meta charset="utf-8">
<title>W18I consistency reference {run_id}</title>
<style>body{{font-family:sans-serif;margin:1.5em}}section{{margin:1.25em 0}}
table{{border-collapse:collapse}}td,th{{border:1px solid #ccc;padding:4px 8px;vertical-align:top}}
.pass{{color:#080}}.fail{{color:#b00}}pre{{white-space:pre-wrap;max-width:110ch}}</style>
</head><body>
<h1>W18I consistency reference {run_id}</h1>
<p>mode: <b>{_esc(mode)}</b> | model: <b>{_esc(args.model)}</b> |
image_api_call_count: <b>{image_api_call_count}</b> |
image_generation_count: <b>{image_generation_count}</b></p>
<section><h2>style contract</h2><pre>{_esc(_MODEST_STYLE_PREAMBLE)}</pre></section>
<section><h2>reference graph rule</h2>
<p>Only L05B02 uses the base FP. Later BGs use at most two prior visual BG
references and do not use the FP directly.</p></section>
{''.join(cards)}
<section><h2>invariants</h2><table><tr><th>invariant</th><th>status</th><th>detail</th></tr>{inv_rows}</table></section>
<details><summary>consistency_reference_plan.json</summary><pre>{_esc(json.dumps(plan, ensure_ascii=False, indent=2))}</pre></details>
</body></html>"""
    )
    run_meta["outputs"].append("run_meta.json")
    (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())