#!/usr/bin/env python
"""Issue 2 probe client: reference-conditioned requests through the API.

Functional pass: one-, two-, and four-reference fixtures at seed 2301.
Measured pass: the four-reference fixture at seeds 2301-2304.

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.

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

from __future__ import annotations

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

import requests

W = pathlib.Path.home() / "Projects/vllm-omni-val/issue2"
SUITE = (
    pathlib.Path.home()
    / ".local/share/multi-model-helper-server/data/evaluation/media/image/references-suite-2"
)
STEPS = 20
CFG = 1.0
FUNCTIONAL_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
)
MEASURED_FIXTURE = "edit-compose-four"
MEASURED_SEEDS = (2301, 2302, 2303, 2304)


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=1200
        )
        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["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("--functional", action="store_true")
    parser.add_argument("--measured", action="store_true")
    parser.add_argument(
        "--fixture",
        default=MEASURED_FIXTURE,
        help="fixture id for --measured (default: edit-compose-four)",
    )
    parser.add_argument(
        "--seeds",
        type=lambda s: tuple(int(x) for x in s.split(",")),
        default=MEASURED_SEEDS,
        help="comma-separated seeds for --measured",
    )
    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.functional:
        for fixture_id in FUNCTIONAL_FIXTURES:
            plan.append(
                (fixture_id, args.out / "functional" / f"{fixture_id}.png", 2301)
            )
    if args.measured:
        for seed in args.seeds:
            plan.append(
                (
                    args.fixture,
                    args.out / "measured" / 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_parts = []
    if args.functional:
        suffix_parts.append("functional")
    if args.measured:
        suffix_parts.append(f"measured-{args.fixture}")
    suffix = "-".join(suffix_parts)
    (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())