from __future__ import annotations

import argparse
import json
import statistics
import threading
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed


def main() -> None:
    parser = argparse.ArgumentParser(description="Run a lightweight chat SSE load test.")
    parser.add_argument("--base-url", default="http://127.0.0.1:8100/api")
    parser.add_argument("--phone", required=True)
    parser.add_argument("--code", default="123456")
    parser.add_argument("--message", default="请用一句话回复：压测验证")
    parser.add_argument("--concurrency", type=int, default=5)
    parser.add_argument("--requests", type=int, default=20)
    parser.add_argument("--timeout", type=int, default=90)
    args = parser.parse_args()

    token = login(args.base_url, args.phone, args.code, args.timeout)
    lock = threading.Lock()
    results: list[dict] = []

    started_at = time.perf_counter()
    with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
        futures = [
            executor.submit(run_one, args.base_url, token, args.message, args.timeout)
            for _ in range(args.requests)
        ]
        for future in as_completed(futures):
            result = future.result()
            with lock:
                results.append(result)
                print(json.dumps(result, ensure_ascii=False))

    elapsed = time.perf_counter() - started_at
    print_summary(results, elapsed)


def login(base_url: str, phone: str, code: str, timeout: int) -> str:
    body = post_json(f"{base_url}/auth/login", {"phone": phone, "code": code}, timeout=timeout)
    return body["data"]["token"]


def run_one(base_url: str, token: str, message: str, timeout: int) -> dict:
    started_at = time.perf_counter()
    try:
        session_body = post_json(
            f"{base_url}/chat/session",
            {"title": "压测会话"},
            token=token,
            timeout=timeout,
        )
        session_id = session_body["data"]["sessionId"]
        first_event_ms = None
        content_events = 0
        error_message = ""

        request = urllib.request.Request(
            f"{base_url}/chat/completions",
            data=json.dumps({"sessionId": session_id, "message": message}).encode("utf-8"),
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {token}",
            },
        )
        with urllib.request.urlopen(request, timeout=timeout) as response:
            for raw_line in response:
                line = raw_line.decode("utf-8").strip()
                if not line.startswith("data:"):
                    continue
                if first_event_ms is None:
                    first_event_ms = elapsed_ms(started_at)
                if line == "data: [DONE]":
                    break
                payload = json.loads(line[5:].strip())
                if payload.get("type") == "content":
                    content_events += 1
                if payload.get("type") == "error":
                    error_message = payload.get("message", "unknown error")

        return {
            "ok": not error_message and content_events > 0,
            "sessionId": session_id,
            "firstEventMs": first_event_ms,
            "totalMs": elapsed_ms(started_at),
            "contentEvents": content_events,
            "error": error_message,
        }
    except (urllib.error.URLError, TimeoutError, KeyError, ValueError) as exc:
        return {
            "ok": False,
            "sessionId": None,
            "firstEventMs": None,
            "totalMs": elapsed_ms(started_at),
            "contentEvents": 0,
            "error": str(exc),
        }


def post_json(url: str, payload: dict, *, token: str | None = None, timeout: int) -> dict:
    headers = {"Content-Type": "application/json"}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    request = urllib.request.Request(url, data=json.dumps(payload).encode("utf-8"), headers=headers)
    with urllib.request.urlopen(request, timeout=timeout) as response:
        return json.loads(response.read().decode("utf-8"))


def print_summary(results: list[dict], elapsed: float) -> None:
    ok_results = [result for result in results if result["ok"]]
    failed_results = [result for result in results if not result["ok"]]
    first_event_values = [result["firstEventMs"] for result in ok_results if result["firstEventMs"] is not None]
    total_values = [result["totalMs"] for result in ok_results]
    summary = {
        "totalRequests": len(results),
        "success": len(ok_results),
        "failed": len(failed_results),
        "elapsedSeconds": round(elapsed, 2),
        "firstEventMsAvg": round(statistics.mean(first_event_values), 2) if first_event_values else None,
        "firstEventMsP95": percentile(first_event_values, 95),
        "totalMsAvg": round(statistics.mean(total_values), 2) if total_values else None,
        "totalMsP95": percentile(total_values, 95),
    }
    print("SUMMARY " + json.dumps(summary, ensure_ascii=False))


def percentile(values: list[int], percentile_value: int) -> int | None:
    if not values:
        return None
    sorted_values = sorted(values)
    index = min(len(sorted_values) - 1, round((percentile_value / 100) * (len(sorted_values) - 1)))
    return sorted_values[index]


def elapsed_ms(started_at: float) -> int:
    return int((time.perf_counter() - started_at) * 1000)


if __name__ == "__main__":
    main()