import argparse
import concurrent.futures
import json
import math
import os
import re
import subprocess
import threading
import time
import urllib.error
import urllib.request
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path


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


def percentile(values, pct):
    values = sorted(v for v in values if v is not None)
    if not values:
        return None
    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 make_tools(count):
    tools = []
    for i in range(count):
        tools.append({
            "type": "function",
            "function": {
                "name": f"tool_{i:02d}_exec",
                "description": "Run a deterministic diagnostic action.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "command": {"type": "string"},
                        "path": {"type": "string"},
                    },
                    "required": ["command"],
                },
            },
        })
    return tools


def make_request(args, idx):
    if args.prompt_mode == "short":
        user = (
            "Do not explain. Output a comma-separated sequence of four digit "
            "numbers starting at 0001. Continue until the token limit stops you."
        )
        messages = [{"role": "user", "content": user}]
    elif args.prompt_mode == "tool":
        messages = [
            {"role": "system", "content": "You are a coding agent. Return concise tool-call-like JSON text."},
            {"role": "user", "content": "Create a shell command to list Python files and print the answer as JSON."},
        ]
    else:
        raise ValueError(f"unknown prompt_mode: {args.prompt_mode}")

    payload = {
        "model": args.model,
        "messages": messages,
        "max_tokens": args.max_tokens,
        "temperature": 0,
        "stream": True,
        "stream_options": {"include_usage": True},
    }
    if args.with_tools:
        payload["tools"] = make_tools(args.tool_count)
        payload["tool_choice"] = "auto"
    return payload


def post_stream(url, payload, timeout):
    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_at = None
    usage = {}
    chars = 0
    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 {}
                    parts = [
                        delta.get("content") or "",
                        delta.get("reasoning_content") or "",
                    ]
                    for tool_call in delta.get("tool_calls") or []:
                        fn = tool_call.get("function") or {}
                        parts.append(fn.get("name") or "")
                        parts.append(fn.get("arguments") or "")
                    added = sum(len(p) for p in parts)
                    if added and first_at is None:
                        first_at = time.perf_counter()
                    chars += added
        elapsed = time.perf_counter() - start
        ttft = first_at - start if first_at is not None else None
        details = usage.get("prompt_tokens_details") or {}
        completion_tokens = int(usage.get("completion_tokens") or 0)
        decode_sec = elapsed - ttft if ttft is not None else elapsed
        return RequestResult(
            ok=True,
            elapsed_sec=elapsed,
            ttft_sec=ttft,
            completion_tokens=completion_tokens,
            prompt_tokens=int(usage.get("prompt_tokens") or 0),
            cached_tokens=int(details.get("cached_tokens") or 0),
            reasoning_tokens=int(usage.get("reasoning_tokens") or 0),
            output_tps=completion_tokens / decode_sec if completion_tokens and decode_sec > 0 else None,
            chars=chars,
        )
    except urllib.error.HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")[:1000]
        return RequestResult(False, time.perf_counter() - start, None, 0, 0, 0, 0, None, chars, f"HTTP {exc.code}: {body}")
    except Exception as exc:
        return RequestResult(False, time.perf_counter() - start, None, 0, 0, 0, 0, None, chars, f"{type(exc).__name__}: {exc}")


def parse_ixsmi(raw):
    samples = []
    for line in raw.splitlines():
        m = re.search(r"\|\s*\d+%\s+\d+C\s+\S+\s+(\d+)W\s*/\s*(\d+)W\s*\|\s*(\d+)MiB\s*/\s*(\d+)MiB\s*\|\s*(\d+)%", line)
        if m:
            samples.append({
                "power_w": int(m.group(1)),
                "power_cap_w": int(m.group(2)),
                "mem_mib": int(m.group(3)),
                "mem_total_mib": int(m.group(4)),
                "util_pct": int(m.group(5)),
            })
    return samples


def monitor_ixsmi(stop_event, out_path, interval):
    records = []
    env = dict(os.environ)
    env["LD_LIBRARY_PATH"] = ":".join([
        "/usr/local/corex/lib64",
        "/usr/local/corex/lib",
        "/usr/local/iluvatar/lib64",
        env.get("LD_LIBRARY_PATH", ""),
    ])
    candidates = [
        "/usr/local/corex/bin/ixsmi",
        "/usr/local/iluvatar/bin/ixsmi",
        "ixsmi",
    ]
    while not stop_event.is_set():
        ts = datetime.now().isoformat(timespec="seconds")
        try:
            last_error = None
            raw = None
            for binary in candidates:
                try:
                    raw = subprocess.check_output([binary], text=True, stderr=subprocess.STDOUT, timeout=10, env=env)
                    break
                except Exception as exc:
                    last_error = exc
            if raw is None:
                raise last_error or RuntimeError("ixsmi not found")
            records.append({"ts": ts, "ok": True, "raw": raw, "parsed": parse_ixsmi(raw)})
        except Exception as exc:
            records.append({"ts": ts, "ok": False, "error": repr(exc)})
        stop_event.wait(interval)
    Path(out_path).write_text(json.dumps(records, ensure_ascii=False, indent=2), encoding="utf-8")


def summarize_monitor(path):
    if not path or not Path(path).exists():
        return None
    records = json.loads(Path(path).read_text(encoding="utf-8"))
    util = []
    mem = []
    power = []
    for rec in records:
        for gpu in rec.get("parsed") or []:
            util.append(gpu["util_pct"])
            mem.append(gpu["mem_mib"])
            power.append(gpu["power_w"])
    if not util:
        return {"records": len(records), "parsed_samples": 0}
    return {
        "records": len(records),
        "parsed_samples": len(util),
        "avg_gpu_util_pct": sum(util) / len(util),
        "max_gpu_util_pct": max(util),
        "avg_mem_mib": sum(mem) / len(mem),
        "max_mem_mib": max(mem),
        "avg_power_w": sum(power) / len(power),
        "max_power_w": max(power),
    }


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("--label", required=True)
    parser.add_argument("--prompt-mode", choices=["short", "tool"], default="short")
    parser.add_argument("--with-tools", action="store_true")
    parser.add_argument("--tool-count", type=int, default=16)
    parser.add_argument("--concurrency", type=int, default=1)
    parser.add_argument("--requests", type=int, default=4)
    parser.add_argument("--max-tokens", type=int, default=256)
    parser.add_argument("--timeout", type=int, default=900)
    parser.add_argument("--monitor-out")
    parser.add_argument("--monitor-interval", type=float, default=1.0)
    parser.add_argument("--out", required=True)
    args = parser.parse_args()

    stop_event = threading.Event()
    monitor_thread = None
    if args.monitor_out:
        monitor_thread = threading.Thread(
            target=monitor_ixsmi,
            args=(stop_event, args.monitor_out, args.monitor_interval),
            daemon=True,
        )
        monitor_thread.start()

    started = time.perf_counter()
    results = []
    try:
        with concurrent.futures.ThreadPoolExecutor(max_workers=args.concurrency) as pool:
            futures = [
                pool.submit(post_stream, args.url, make_request(args, i), args.timeout)
                for i in range(args.requests)
            ]
            for fut in concurrent.futures.as_completed(futures):
                result = fut.result()
                results.append(result)
                print(f"done {len(results)}/{args.requests} ok={result.ok} tps={result.output_tps}", flush=True)
    finally:
        stop_event.set()
        if monitor_thread:
            monitor_thread.join(timeout=15)

    wall = time.perf_counter() - started
    ok = [r for r in results if r.ok]
    tps_values = [r.output_tps for r in ok if r.output_tps is not None]
    ttft_values = [r.ttft_sec for r in ok if r.ttft_sec is not None]
    completion = sum(r.completion_tokens for r in ok)
    prompt = sum(r.prompt_tokens for r in ok)
    cached = sum(r.cached_tokens for r in ok)
    summary = {
        "created_at": datetime.now().isoformat(timespec="seconds"),
        "label": args.label,
        "url": args.url,
        "model": args.model,
        "prompt_mode": args.prompt_mode,
        "with_tools": args.with_tools,
        "tool_count": args.tool_count if args.with_tools else 0,
        "concurrency": args.concurrency,
        "requests": args.requests,
        "max_tokens": args.max_tokens,
        "wall_sec": wall,
        "success_rate": len(ok) / len(results) if results else 0,
        "ttft_p50_sec": percentile(ttft_values, 50),
        "ttft_p90_sec": percentile(ttft_values, 90),
        "output_tps_p10_per_request": percentile(tps_values, 10),
        "output_tps_p50_per_request": percentile(tps_values, 50),
        "aggregate_output_tps": completion / wall if wall > 0 else 0,
        "prompt_tokens": prompt,
        "cached_tokens": cached,
        "completion_tokens": completion,
        "reasoning_tokens": sum(r.reasoning_tokens for r in ok),
        "chars": sum(r.chars for r in ok),
        "monitor": summarize_monitor(args.monitor_out) if args.monitor_out else None,
        "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()