#!/usr/bin/env python
"""Issue 5 probe client: reference-conditioned requests against the int8
weight-only compressed-tensors artifact, with the structure gate applied to
every response.

Matrix pass (the issue's capacity matrix): one-, two-, three-, and
four-reference fixtures at seeds 2301-2304 — wall time and per-image
structure metrics land in the record for every request. Every request
mirrors the sealed Diffusers baseline attempt 2026-10-07T170729Z-q3-4r
(same frozen references, same rendered prompt, 20 steps, true_cfg_scale
1.0, fixture resolution, no negative prompt), submitted one at a time
through POST /v1/images/edits.

The structure gate (the lesson from issue 1's FP8 misread): a run only
counts when the output is structurally sound — fully opaque RGBA and
non-flat RGB. Broken FP8 output measured ~96.5-97% transparent with
near-flat grey RGB; the BF16 path is 100% opaque. The gate numbers are
recorded on every request regardless; `gate.py` summarizes them.

Usage: probe.py --out RUN_DIR [--matrix] [--host H:P]
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import io
import json
import pathlib
import time

import numpy as np
import requests
from PIL import Image

W = pathlib.Path.home() / "Projects/vllm-omni-val/issue5"
SUITE = (
    pathlib.Path.home()
    / ".local/share/multi-model-helper-server/data/evaluation/media/image/references-suite-2"
)
STEPS = 20
CFG = 1.0
MATRIX_FIXTURES = (
    "edit-expression-ysolde",  # 1 reference
    "edit-compose-pair",        # 2 references
    "edit-compose-three",       # 3 references (derived; see fixtures.json)
    "edit-compose-four",        # 4 references
)
MATRIX_SEEDS = (2301, 2302, 2303, 2304)


def structure_metrics(png: bytes) -> dict:
    """Opacity + non-flatness numbers for the structure gate."""
    image = Image.open(io.BytesIO(png))
    array = np.asarray(image.convert("RGBA"))
    alpha = array[..., 3]
    rgb = array[..., :3].astype(np.float64)
    return {
        "opaque_fraction": round(float((alpha == 255).mean()), 4),
        "alpha_mean": round(float(alpha.mean()), 1),
        "rgb_std": [round(float(rgb[..., i].std()), 1) for i in range(3)],
    }


def submit(host: str, fixture: dict, seed: int) -> tuple[dict, bytes | None]:
    files = []
    handles = []
    for ref in fixture["references"]:
        path = SUITE / f"{ref}.png"
        handle = open(path, "rb")
        handles.append(handle)
        files.append(("image[]", (path.name, handle, "image/png")))
    data = {
        "model": "Qwen/Qwen-Image-2.1",
        "prompt": fixture["prompt"],
        "seed": str(seed),
        "num_inference_steps": str(STEPS),
        "true_cfg_scale": str(CFG),
        "size": f"{fixture['width']}x{fixture['height']}",
        "n": "1",
        "response_format": "b64_json",
        "output_format": "png",
        "return_stage_metrics": "true",
    }
    rec: dict = {
        "fixture_id": None,
        "seed": seed,
        "n_references": len(fixture["references"]),
        "steps": STEPS,
        "true_cfg_scale": CFG,
        "size": f"{fixture['width']}x{fixture['height']}",
    }
    t0 = time.time()
    try:
        response = requests.post(
            f"http://{host}/v1/images/edits", data=data, files=files, timeout=1800
        )
        rec["wall_s"] = round(time.time() - t0, 2)
        rec["http"] = response.status_code
        if response.status_code == 200:
            body = response.json()
            png = base64.b64decode(body["data"][0]["b64_json"])
            rec["bytes"] = len(png)
            rec["sha256"] = hashlib.sha256(png).hexdigest()
            rec["structure"] = structure_metrics(png)
            rec["usage"] = body.get("usage")
            for key in ("stage_metrics", "metrics", "timings"):
                if body.get(key):
                    rec[key] = body[key]
            return rec, png
        rec["error"] = response.text[:2000]
        return rec, None
    except Exception as exc:  # noqa: BLE001 - recorded as the probe result
        rec["wall_s"] = round(time.time() - t0, 2)
        rec["error"] = f"{type(exc).__name__}: {exc}"
        return rec, None
    finally:
        for handle in handles:
            handle.close()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--out", required=True, type=pathlib.Path)
    parser.add_argument("--host", default="127.0.0.1:8091")
    parser.add_argument("--matrix", action="store_true")
    parser.add_argument(
        "--fixture",
        default="edit-compose-four",
        help="single fixture id for --single",
    )
    parser.add_argument("--single", action="store_true", help="one fixture, all seeds")
    parser.add_argument(
        "--seeds",
        type=lambda s: tuple(int(x) for x in s.split(",")),
        default=MATRIX_SEEDS,
    )
    args = parser.parse_args()

    fixtures = json.loads((W / "fixtures.json").read_text())
    args.out.mkdir(parents=True, exist_ok=True)

    plan: list[tuple[str, pathlib.Path, int]] = []
    if args.matrix:
        for fixture_id in MATRIX_FIXTURES:
            for seed in args.seeds:
                plan.append(
                    (
                        fixture_id,
                        args.out / "matrix" / f"{fixture_id}.s{seed}.png",
                        seed,
                    )
                )
    if args.single:
        for seed in args.seeds:
            plan.append(
                (args.fixture, args.out / "matrix" / f"{args.fixture}.s{seed}.png", seed)
            )

    records = []
    for fixture_id, out_path, seed in plan:
        rec, png = submit(args.host, fixtures[fixture_id], seed)
        rec["fixture_id"] = fixture_id
        records.append(rec)
        print(json.dumps(rec), flush=True)
        if png is not None:
            out_path.parent.mkdir(parents=True, exist_ok=True)
            out_path.write_bytes(png)
        elif rec.get("http") == 200:
            print(f"warning: 200 but no image for {fixture_id} seed {seed}", flush=True)

    suffix = "matrix" if args.matrix else f"single-{args.fixture}"
    (args.out / f"{suffix}.json").write_text(json.dumps(records, indent=1) + "\n")
    ok = all(r.get("http") == 200 and "sha256" in r for r in records)
    print(("PASS" if ok else "FAIL") + f" {len(records)} requests", flush=True)
    return 0 if ok else 1


if __name__ == "__main__":
    raise SystemExit(main())