import argparse
import concurrent.futures
import json
import math
import statistics
import time
import urllib.error
import urllib.request
from dataclasses import dataclass, asdict
from datetime import datetime
from pathlib import Path


@dataclass
class RequestResult:
    ok: bool
    elapsed_sec: float
    ttft_sec: float | None
    prompt_tokens: int
    cached_tokens: int
    completion_tokens: int
    reasoning_tokens: int
    output_tps: float | None
    error: str | None = None


def percentile(values, pct):
    if not values:
        return None
    values = sorted(values)
    if len(values) == 1:
        return values[0]
    pos = (len(values) - 1) * pct / 100.0
    lo = math.floor(pos)
    hi = math.ceil(pos)
    if lo == hi:
        return values[lo]
    return values[lo] * (hi - pos) + values[hi] * (pos - lo)


def load_dataset(path):
    text = Path(path).read_text(encoding="utf-8")
    if path.endswith(".jsonl"):
        data = [json.loads(line) for line in text.splitlines() if line.strip()]
    else:
        data = json.loads(text)

    # Official-like format: each item is already an OpenAI chat completion
    # request with messages/tools/tool_choice/stream/etc.
    if data and isinstance(data[0], dict) and "messages" in data[0]:
        return data

    requests = []
    for item in data:
        system_prompt = item.get("system_prompt") or "You are a helpful assistant."
        history = [{"role": "system", "content": system_prompt}]
        for question in item.get("user_questions", []):
            messages = history + [{"role": "user", "content": question}]
            requests.append({"messages": messages})
            # Synthetic assistant placeholder keeps later prompts multi-turn.
            history = messages + [{"role": "assistant", "content": "好的，我们继续。"}]
    return requests


def post_stream(url, model, request_item, max_tokens, timeout):
    allowed_fields = {
        "messages",
        "tools",
        "tool_choice",
        "model",
        "max_tokens",
        "temperature",
        "top_p",
        "stop",
        "presence_penalty",
        "frequency_penalty",
        "n",
        "response_format",
    }
    payload = {k: v for k, v in dict(request_item).items() if k in allowed_fields}
    payload["model"] = payload.get("model") or model
    # The benchmark CLI controls output length, even if the synthetic dataset
    # stores a larger official-like max_tokens value.
    payload["max_tokens"] = max_tokens
    payload["temperature"] = payload.get("temperature", 0)
    payload["stream"] = True
    payload["stream_options"] = {"include_usage": True}
    req = urllib.request.Request(
        url.rstrip("/") + "/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 = {}
    try:
        with urllib.request.urlopen(req, timeout=timeout) 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"]
                for choice in obj.get("choices") or []:
                    delta = choice.get("delta") or {}
                    text = delta.get("content") or delta.get("reasoning_content") or ""
                    if text and first_token_at is None:
                        first_token_at = time.perf_counter()
        elapsed = time.perf_counter() - start
        ttft = first_token_at - start if first_token_at else None
        completion_tokens = int(usage.get("completion_tokens") or 0)
        prompt_tokens = int(usage.get("prompt_tokens") or 0)
        reasoning_tokens = int(usage.get("reasoning_tokens") or 0)
        details = usage.get("prompt_tokens_details") or {}
        cached_tokens = int(details.get("cached_tokens") or 0)
        decode_sec = elapsed - ttft if ttft is not None else elapsed
        output_tps = completion_tokens / decode_sec if completion_tokens and decode_sec > 0 else None
        return RequestResult(
            ok=True,
            elapsed_sec=elapsed,
            ttft_sec=ttft,
            prompt_tokens=prompt_tokens,
            cached_tokens=cached_tokens,
            completion_tokens=completion_tokens,
            reasoning_tokens=reasoning_tokens,
            output_tps=output_tps,
        )
    except urllib.error.HTTPError as exc:
        body = ""
        try:
            body = exc.read().decode("utf-8", errors="replace")
        except Exception:
            pass
        elapsed = time.perf_counter() - start
        return RequestResult(
            ok=False,
            elapsed_sec=elapsed,
            ttft_sec=None,
            prompt_tokens=0,
            cached_tokens=0,
            completion_tokens=0,
            reasoning_tokens=0,
            output_tps=None,
            error=f"HTTPError {exc.code}: {body[:1000]}",
        )
    except Exception as exc:
        elapsed = time.perf_counter() - start
        return RequestResult(
            ok=False,
            elapsed_sec=elapsed,
            ttft_sec=None,
            prompt_tokens=0,
            cached_tokens=0,
            completion_tokens=0,
            reasoning_tokens=0,
            output_tps=None,
            error=f"{type(exc).__name__}: {exc}",
        )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--url", default="http://127.0.0.1:1111")
    parser.add_argument("--model", default="llm")
    parser.add_argument("--dataset", required=True)
    parser.add_argument("--concurrency", type=int, default=1)
    parser.add_argument("--max-requests", type=int, default=20)
    parser.add_argument("--max-tokens", type=int, default=128)
    parser.add_argument("--timeout", type=int, default=600)
    parser.add_argument("--out", default="/root/work/logs/formal_perf_result.json")
    args = parser.parse_args()

    reqs = load_dataset(args.dataset)
    if not reqs:
        raise SystemExit("empty dataset")
    scheduled = [reqs[i % len(reqs)] for i in range(args.max_requests)]

    wall_start = time.perf_counter()
    results = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.concurrency) as pool:
        futures = [
            pool.submit(post_stream, args.url, args.model, request_item, args.max_tokens, args.timeout)
            for request_item in scheduled
        ]
        for fut in concurrent.futures.as_completed(futures):
            results.append(fut.result())
            print(f"done {len(results)}/{len(scheduled)} ok={results[-1].ok}", flush=True)
    wall_sec = time.perf_counter() - wall_start

    ok_results = [r for r in results if r.ok]
    success_rate = len(ok_results) / len(results) if results else 0.0
    ttfts = [r.ttft_sec for r in ok_results if r.ttft_sec is not None]
    output_tps_values = [r.output_tps for r in ok_results if r.output_tps is not None]

    prompt_tokens = sum(r.prompt_tokens for r in ok_results)
    cached_tokens = sum(r.cached_tokens for r in ok_results)
    input_tokens_uncached = max(prompt_tokens - cached_tokens, 0)
    output_tokens = sum(r.completion_tokens for r in ok_results)

    aggregate_output_tps = output_tokens / wall_sec if wall_sec > 0 else 0.0
    aggregate_input_tps = input_tokens_uncached / wall_sec if wall_sec > 0 else 0.0
    aggregate_cache_tps = cached_tokens / wall_sec if wall_sec > 0 else 0.0
    weighted = aggregate_output_tps * 16.796 + aggregate_input_tps * 2.799 + aggregate_cache_tps * 0.56
    cache_hit_rate = cached_tokens / prompt_tokens if prompt_tokens else 0.0

    summary = {
        "created_at": datetime.now().isoformat(timespec="seconds"),
        "url": args.url,
        "model": args.model,
        "dataset": args.dataset,
        "concurrency": args.concurrency,
        "max_requests": args.max_requests,
        "max_tokens": args.max_tokens,
        "wall_sec": wall_sec,
        "success_rate": success_rate,
        "ttft_p90_sec": percentile(ttfts, 90),
        "output_tps_p10_per_request": percentile(output_tps_values, 10),
        "aggregate_output_tps": aggregate_output_tps,
        "aggregate_input_tps_uncached": aggregate_input_tps,
        "aggregate_cache_tps": aggregate_cache_tps,
        "cache_hit_rate": cache_hit_rate,
        "weighted_token_throughput": weighted,
        "totals": {
            "requests": len(results),
            "success": len(ok_results),
            "prompt_tokens": prompt_tokens,
            "input_tokens_uncached": input_tokens_uncached,
            "cached_tokens": cached_tokens,
            "completion_tokens": output_tokens,
            "reasoning_tokens": sum(r.reasoning_tokens for r in ok_results),
        },
        "targets": {
            "output_tps_p10_per_request_gte_20": (percentile(output_tps_values, 10) or 0) >= 20,
            "ttft_p90_lte_5": (percentile(ttfts, 90) or 999) <= 5,
            "cache_hit_rate_gte_50pct": cache_hit_rate >= 0.5,
            "success_rate_gte_99pct": success_rate >= 0.99,
            "weighted_token_throughput_gte_8000": weighted >= 8000,
        },
        "results": [asdict(r) for r in results],
    }

    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    Path(args.out).write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps({k: v for k, v in summary.items() if k != "results"}, ensure_ascii=False, indent=2))
    print("RESULT_FILE", args.out)


if __name__ == "__main__":
    main()