"""Reproduce the real-hardware API checks. Run under an external timeout."""
import base64
import argparse
from datetime import datetime,timezone
import hashlib
import json
import os
from pathlib import Path
import subprocess
import sys
import time
import traceback
from urllib.request import Request,urlopen
from urllib.error import HTTPError,URLError

PACKAGE=Path(__file__).resolve().parent.parent
parser=argparse.ArgumentParser()
parser.add_argument("--backend",choices=("native","reference","quark-hip"),default="native")
parser.add_argument("--benchmark-only",action="store_true")
parser.add_argument("--model-dir",type=Path,required=True)
parser.add_argument("--output",type=Path,required=True)
parser.add_argument("--quark-extension",type=Path)
options=parser.parse_args()
if options.backend=="quark-hip" and options.quark_extension is None:
    parser.error("quark-hip requires --quark-extension from the pinned Quark wheel")
out=options.output.resolve()
out.mkdir(parents=True,exist_ok=False)
env=dict(os.environ)
started=time.perf_counter()
report=dict(status="running",requests=[],checkpoint="balanced",backend=options.backend,transport="local HTTP JSON/base64 PNG")
server_log=(out/"server.log").open("w")
command=[str(PACKAGE/"launch.sh"),"--model-dir",str(options.model_dir.resolve()),
         "--backend",options.backend,"serve","--port","8191"]
if options.backend=="quark-hip":
    command=[sys.executable,"-u",str(PACKAGE/"tools/quark_baseline_server.py"),
             "--model-dir",str(options.model_dir.resolve()),"--extension",str(options.quark_extension.resolve())]
process=subprocess.Popen(command,env=env,stdout=server_log,stderr=subprocess.STDOUT)

def save():
    (out/"result.json").write_text(json.dumps(report,indent=2)+"\n")

def post(name,body,edit=False):
    begin=time.perf_counter()
    path="/v1/images/edits" if edit else "/v1/images/generations"
    request=Request("http://127.0.0.1:8191"+path,data=json.dumps(body).encode(),headers={"Content-Type":"application/json"})
    with urlopen(request,timeout=340) as response:
        data=json.load(response)
    png=base64.b64decode(data["data"][0]["b64_json"],validate=True)
    (out/(name+".png")).write_bytes(png)
    row=dict(name=name,http_request_seconds=time.perf_counter()-begin,process_start_to_response_seconds=time.perf_counter()-started,
             png_sha256=hashlib.sha256(png).hexdigest(),**data["metrics"])
    report["requests"].append(row)
    save()
    print("REQUEST COMPLETE",json.dumps(row),flush=True)
    return png,row

print("OUTPUT:",out,flush=True)
try:
    for attempt in range(150):
        if process.poll() is not None:
            raise RuntimeError("Server exited during startup; see server.log")
        try:
            with urlopen("http://127.0.0.1:8191/health",timeout=1) as f:
                report["health"]=json.load(f)
            break
        except (URLError,TimeoutError):
            time.sleep(1)
    else:
        raise RuntimeError("Server startup exceeded 150 seconds")
    report["startup_to_ready_seconds"]=time.perf_counter()-started
    with urlopen("http://127.0.0.1:8191/v1/models",timeout=2) as f:
        report["models"]=json.load(f)
    for body in (dict(prompt="test",steps=1),dict(prompt="test",size="4096x4096"),dict(prompt="test",n=2)):
        request=Request("http://127.0.0.1:8191/v1/images/generations",data=json.dumps(body).encode(),headers={"Content-Type":"application/json"})
        try:
            urlopen(request,timeout=3)
            raise AssertionError("Invalid profile accepted")
        except HTTPError as error:
            assert error.code==400
    report["invalid_profile_rejection"]=True
    prompt='A neon shop sign that reads "QWEN IMAGE 2.1", rainy night, reflections on wet pavement'
    large=dict(prompt=prompt,size="2048x2048",seed=42)
    first,first_row=post("neon-2048-first",large)
    warm,warm_row=post("neon-2048-warm",large)
    assert first==warm,"Repeated full-resolution output changed"
    if options.benchmark_only:
        report.update(status="passed_API_benchmark",repeat_2048_png_identical=True)
        save()
        raise SystemExit(0)
    a_body=dict(prompt="A blue ceramic teapot on a wooden table, soft daylight, studio photograph",size="1024x1024",seed=51)
    a,a_row=post("cache-a-first",a_body)
    rgba,rgba_row=post("rgba-1024",dict(prompt="A cute cartoon dragon sticker.",mode="rgba",size="1024x1024",seed=52))
    assert rgba_row["alpha_extrema"][0]<64 and rgba_row["alpha_extrema"][1]>192,"No meaningful transparency"
    edit,edit_row=post("edit-1024",dict(prompt="Change the blue teapot to bright red. Keep the table and composition.",
        image_b64=base64.b64encode(a).decode(),size="1024x1024",seed=53),edit=True)
    repeat,repeat_row=post("cache-a-repeat",a_body)
    assert a==repeat,"A-B-A output changed after editing and RGBA requests"
    report.update(status="passed_API_hardware_screen",repeat_2048_png_identical=True,aba_png_identical=True,
                  limits="batch one; text-to-image 1024/2048 square; RGBA 1024 square tested; edit one input at 1024 square; 40 steps; CFG1; non-default stream; no full-pipeline graph capture")
except Exception:
    report["status"]="failed"
    report["traceback"]=traceback.format_exc()
    raise
finally:
    process.terminate()
    try:
        process.wait(timeout=15)
    except subprocess.TimeoutExpired:
        process.kill()
        process.wait(timeout=5)
    server_log.close()
    report["elapsed_seconds"]=time.perf_counter()-started
    save()
print("API VALIDATION COMPLETE:",out/"result.json",flush=True)