#!/usr/bin/env python3
"""Prove a selected context is real, then verify reset after the long prompt."""
import argparse
import json
import urllib.request

p = argparse.ArgumentParser()
p.add_argument("--url", required=True)
p.add_argument("--digits", type=int, default=4200)
args = p.parse_args()
opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))


def chat(text, count):
    body = {"model": "mindnano-ling3-tiny", "messages": [{"role": "user", "content": text}],
            "max_tokens": count, "temperature": 0}
    req = urllib.request.Request(args.url + "/v1/chat/completions", json.dumps(body).encode(),
                                 {"Content-Type": "application/json"})
    with opener.open(req, timeout=300) as response:
        return json.load(response)


a = chat("请只回答一个中文问候。", 16)
long = chat("这是测试编号，请只回答收到：" + "1" * args.digits, 1)
assert long["usage"]["prompt_tokens"] > 4096, long["usage"]
b = chat("请只回答一个中文问候。", 16)
assert a["choices"][0]["message"]["content"] == b["choices"][0]["message"]["content"], "long-context reset leaked state"
print(json.dumps({"passed": True, "long_request": long, "reset_reply": b}, ensure_ascii=False, indent=2))