"""Explicit test against an already running v2 engine; performs real inference."""
import argparse
import json
import time
import urllib.error
import urllib.request


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--base", default="http://127.0.0.1:19091")
    args = parser.parse_args()
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))

    def open_request(path, body=None):
        return opener.open(urllib.request.Request(args.base + path,
            data=None if body is None else json.dumps(body).encode(),
            headers={"Content-Type": "application/json"}), timeout=180)

    def request(path, body=None):
        with open_request(path, body) as response:
            return json.load(response)

    def events(response):
        for line in response:
            if line.startswith(b"data:"):
                data = line[5:].strip()
                if data == b"[DONE]":
                    return
                yield json.loads(data)

    def rejected(path, body):
        try:
            request(path, body)
        except urllib.error.HTTPError as exc:
            assert exc.code == 409, exc.read().decode()
            return
        raise AssertionError("stale control accepted")

    health = request("/health")
    assert health["mindnano_integration_version"] >= 2
    body = {"model": "mindnano-ling3-tiny", "messages": [{"role": "user", "content": "讲一个小兔子探索森林的长故事。"}],
            "stream": True, "max_tokens": 128, "temperature": 0, "top_k": 20, "repeat_penalty": 1.05}
    with open_request("/v1/chat/completions", dict(body, flow_control="ack")) as response:
        stream = events(response)
        first = next(event for event in stream if event.get("mindnano_flow"))
        flow = first["mindnano_flow"]
        before = request("/v1/generation/status")
        time.sleep(0.5)
        after = request("/v1/generation/status")
        assert before["waiting_for_ack"] and after["waiting_for_ack"]
        assert before["evaluated_steps"] == after["evaluated_steps"]
        rejected("/v1/flow/ack", dict(flow, request_id="stale"))
        rejected("/v1/flow/ack", dict(flow, sequence=flow["sequence"] + 1))
        assert request("/v1/cancel", {"request_id": "stale"})["cancel_requested"] is False
        request("/v1/flow/ack", flow)
        second = next(event for event in stream if event.get("mindnano_flow"))
        assert second["mindnano_flow"]["sequence"] == flow["sequence"] + 1
        resumed = request("/v1/generation/status")
        assert resumed["evaluated_steps"] > after["evaluated_steps"]
        request("/v1/cancel", {"request_id": flow["request_id"]})
        tail = list(stream)
        assert any(event.get("error", {}).get("code") == "request_canceled" for event in tail)
    print(json.dumps({"flow_ack": "PASS", "parked_steps": after["evaluated_steps"],
                      "resumed_steps": resumed["evaluated_steps"], "stale_controls": "PASS"}), flush=True)

    with open_request("/v1/chat/completions", body) as response:
        stream = events(response)
        first = next(stream)
        request_id = first["id"]
        paused = request("/v1/generation/pause", {"request_id": request_id})
        assert paused["paused"]
        time.sleep(0.5)
        parked = request("/v1/generation/status")
        assert parked["evaluated_steps"] == paused["evaluated_steps"]
        assert parked["paused"]
        rejected("/v1/generation/resume", {"request_id": flow["request_id"]})
        request("/v1/generation/resume", {"request_id": request_id})
        next(event for event in stream if any(c.get("delta", {}).get("content") for c in event.get("choices", [])))
        request("/v1/cancel", {"request_id": request_id})
        list(stream)
    print(json.dumps({"pause_resume": "PASS", "parked_steps": parked["evaluated_steps"]}), flush=True)

    with open_request("/v1/chat/completions", dict(body, max_tokens=8, flow_control="ack")) as response:
        stream = events(response)
        first = next(event for event in stream if event.get("mindnano_flow"))
        # Closing without acknowledging must unblock the inference thread.
    deadline = time.monotonic() + 5
    while request("/v1/generation/status")["active"]:
        assert time.monotonic() < deadline, "disconnect did not cancel parked generation"
        time.sleep(0.05)
    result = request("/v1/chat/completions", dict(body, stream=False, max_tokens=8))
    assert result["usage"]["completion_tokens"] > 0
    assert result["mindnano_metrics"]["sampling"]["top_k"] == 20
    assert result["mindnano_metrics"]["sampling"]["repeat_penalty"] == 1.05
    print(json.dumps({"disconnect_recovery": "PASS", "sampling": result["mindnano_metrics"]["sampling"]}), flush=True)


if __name__ == "__main__":
    main()