#!/usr/bin/env python3
"""Run against an initialized portable service; Python is client-side only."""
import argparse
import concurrent.futures
import http.client
import json
import time
import urllib.error
import urllib.request


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--url", default="http://127.0.0.1:19093")
    args = p.parse_args()
    base = args.url.rstrip("/")

    def call(path, body=None):
        request = urllib.request.Request(base + path, None if body is None else json.dumps(body).encode(),
                                         {"Content-Type": "application/json"})
        with urllib.request.urlopen(request, timeout=120) as r:
            return json.load(r)

    def request(text, **kwargs):
        return {"model": "mindnano-ling3-tiny", "messages": [{"role": "user", "content": text}],
                "max_tokens": 24, "temperature": 0, "cache_prompt": False, **kwargs}

    health = call("/health")
    assert health["status"] == "ok"
    assert call("/v1/models")["data"][0]["context_length"] == health["context_length"]
    for patch, status in [({"model": "bad"}, 404), ({"max_tokens": -1}, 400),
                          ({"tools": []}, 400), ({"messages": []}, 400),
                          ({"messages": [{"role": "user", "content": "1" * (health["context_length"] + 1)}]}, 400)]:
        try:
            call("/v1/chat/completions", request("你好", **patch))
            raise AssertionError("invalid request accepted")
        except urllib.error.HTTPError as e:
            assert e.code == status, (e.code, e.read())
            assert "error" in json.load(e)

    a = call("/v1/chat/completions", request("请用一句话介绍你自己。"))
    b = call("/v1/chat/completions", request("只回答：一加一等于几？"))
    again = call("/v1/chat/completions", request("请用一句话介绍你自己。"))
    assert a["choices"][0]["message"]["content"] == again["choices"][0]["message"]["content"], "request isolation/reset failed"
    for r in (a, b, again):
        assert r["object"] == "chat.completion"
        assert r["usage"]["total_tokens"] == r["usage"]["prompt_tokens"] + r["usage"]["completion_tokens"]
        assert r["mindnano_metrics"]["cached_tokens"] == 0

    r = urllib.request.Request(base + "/v1/chat/completions", json.dumps(
        request("请用一句话介绍你自己。", stream=True, stream_options={"include_usage": True})).encode(),
        {"Content-Type": "application/json"})
    content = ""; done = False; metrics = None; usage = None
    with urllib.request.urlopen(r, timeout=120) as response:
        for raw in response:
            line = raw.decode("utf-8").strip()
            if not line.startswith("data: "):
                continue
            if line == "data: [DONE]":
                done = True; break
            event = json.loads(line[6:])
            assert "error" not in event, event
            for c in event["choices"]:
                content += c["delta"].get("content", "")
            metrics = event.get("mindnano_metrics", metrics)
            usage = event.get("usage") or usage
    assert done and usage and metrics
    assert content == a["choices"][0]["message"]["content"], "SSE text mismatch"

    # Cancel non-stream generation; health/cancel must remain available while NPU is busy.
    with concurrent.futures.ThreadPoolExecutor() as pool:
        future = pool.submit(call, "/v1/chat/completions", request("详细写一篇长篇故事。", max_tokens=1024))
        for _ in range(100):
            if call("/health")["active"]:
                break
            time.sleep(.02)
        else:
            raise AssertionError("generation did not start")
        try:
            call("/v1/chat/completions", request("并发请求"))
            raise AssertionError("concurrent request was not rejected")
        except urllib.error.HTTPError as e:
            assert e.code == 429
        assert call("/v1/cancel", {})["cancel_requested"]
        try:
            future.result(timeout=10)
            raise AssertionError("cancel returned successful completion")
        except urllib.error.HTTPError as e:
            assert e.code == 409
    final = call("/v1/chat/completions", request("请用一句话介绍你自己。"))
    assert final["choices"][0]["message"]["content"] == content, "canceled state leaked"
    print(json.dumps({"passed": True, "context": health["context_length"],
                      "sample": content, "metrics": [a["mindnano_metrics"], b["mindnano_metrics"], metrics]},
                     ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()