import json
import statistics
import sys
import time
import urllib.request
from datetime import datetime
from pathlib import Path


BASE_URL = sys.argv[1] if len(sys.argv) > 1 else "http://127.0.0.1:1111"
OUT_DIR = Path(sys.argv[2]) if len(sys.argv) > 2 else Path("/root/work/logs")
MODEL = sys.argv[3] if len(sys.argv) > 3 else "llm"


def post_json(path, payload, timeout=300):
    req = urllib.request.Request(
        BASE_URL + path,
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return resp.status, resp.read().decode("utf-8")


def health():
    with urllib.request.urlopen(BASE_URL + "/health", timeout=10) as resp:
        return resp.status


def chat_once(prompt, max_tokens=32):
    payload = {
        "model": MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "temperature": 0,
    }
    start = time.perf_counter()
    status, body = post_json("/v1/chat/completions", payload)
    elapsed = time.perf_counter() - start
    obj = json.loads(body)
    usage = obj.get("usage") or {}
    completion_tokens = usage.get("completion_tokens") or 0
    return {
        "status": status,
        "elapsed_sec": elapsed,
        "usage": usage,
        "output_tps": completion_tokens / elapsed if elapsed and completion_tokens else None,
        "content_preview": (obj["choices"][0]["message"].get("content") or "")[:200],
    }


def chat_stream(prompt, max_tokens=64):
    payload = {
        "model": MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "temperature": 0,
        "stream": True,
        "stream_options": {"include_usage": True},
    }
    req = urllib.request.Request(
        BASE_URL + "/v1/chat/completions",
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    start = time.perf_counter()
    first_token_at = None
    usage = None
    pieces = []
    with urllib.request.urlopen(req, timeout=300) as resp:
        for raw in resp:
            line = raw.decode("utf-8", errors="replace").strip()
            if not line or not line.startswith("data:"):
                continue
            data = line[5:].strip()
            if data == "[DONE]":
                break
            obj = json.loads(data)
            if obj.get("usage"):
                usage = obj["usage"]
            choices = obj.get("choices") or []
            if choices:
                delta = choices[0].get("delta") or {}
                text = delta.get("content") or delta.get("reasoning_content") or ""
                if text:
                    if first_token_at is None:
                        first_token_at = time.perf_counter()
                    pieces.append(text)
    elapsed = time.perf_counter() - start
    completion_tokens = (usage or {}).get("completion_tokens") or 0
    gen_time = elapsed - (first_token_at - start) if first_token_at else elapsed
    return {
        "elapsed_sec": elapsed,
        "ttft_sec": (first_token_at - start) if first_token_at else None,
        "usage": usage,
        "output_tps_after_ttft": completion_tokens / gen_time if gen_time and completion_tokens else None,
        "content_preview": "".join(pieces)[:200],
    }


def main():
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    results = {
        "base_url": BASE_URL,
        "model": MODEL,
        "created_at": datetime.now().isoformat(timespec="seconds"),
    }
    results["health_status"] = health()
    results["smoke_nonstream"] = chat_once("你好，请用一句话介绍你自己。", max_tokens=16)

    short_prompt = "请用中文简要说明什么是模型推理服务。"
    stream_runs = [chat_stream(short_prompt, max_tokens=64) for _ in range(3)]
    results["stream_runs"] = stream_runs
    ttfts = [x["ttft_sec"] for x in stream_runs if x["ttft_sec"] is not None]
    tps = [x["output_tps_after_ttft"] for x in stream_runs if x["output_tps_after_ttft"] is not None]
    results["stream_summary"] = {
        "ttft_avg_sec": statistics.mean(ttfts) if ttfts else None,
        "ttft_p90_sec": sorted(ttfts)[int(0.9 * (len(ttfts) - 1))] if ttfts else None,
        "output_tps_avg": statistics.mean(tps) if tps else None,
    }

    repeated = "以下是一段用于测试前缀缓存的公共上下文：" + ("模型部署竞赛关注吞吐、延迟、缓存命中和稳定性。" * 80)
    results["prefix_cache_probe_1"] = chat_once(repeated + "\n请总结一句话。", max_tokens=16)
    results["prefix_cache_probe_2"] = chat_once(repeated + "\n请换一种说法总结一句话。", max_tokens=16)

    out = OUT_DIR / f"baseline_smoke_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    out.write_text(json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(results, ensure_ascii=False, indent=2), flush=True)
    print("RESULT_FILE", out, flush=True)


if __name__ == "__main__":
    main()