#!/usr/bin/env python3
"""Measure direct MLX-LM generation without an HTTP serving layer."""

from __future__ import annotations

import argparse
import json
import math
import platform
import statistics
import time
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

import mlx.core as mx
from benchmark import PROMPTS
from mlx_lm import load, stream_generate
from mlx_lm.sample_utils import make_sampler


def percentile(values: list[float], fraction: float) -> float:
    ordered = sorted(values)
    position = (len(ordered) - 1) * fraction
    lower = math.floor(position)
    upper = math.ceil(position)
    if lower == upper:
        return ordered[lower]
    return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower)


def summarize(rows: list[dict[str, Any]]) -> dict[str, Any]:
    measured = [row for row in rows if row["phase"] == "measured"]
    result: dict[str, Any] = {
        "runtime": "mlx-lm-direct",
        "model": measured[0]["model"],
        "requests": len(measured),
        "successes": sum(row["error"] is None for row in measured),
    }
    for field in ("ttft_s", "total_s", "prompt_tps", "decode_tps", "request_tps"):
        values = [float(row[field]) for row in measured if row[field] is not None]
        result[f"{field}_mean"] = statistics.fmean(values)
        result[f"{field}_p50"] = percentile(values, 0.5)
        result[f"{field}_p95"] = percentile(values, 0.95)
    result["peak_memory_gb"] = max(float(row["peak_memory_gb"]) for row in measured)
    return result


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True)
    parser.add_argument("--label", required=True)
    parser.add_argument("--warmups", type=int, default=2)
    parser.add_argument("--rounds", type=int, default=5)
    parser.add_argument("--max-tokens", type=int, default=128)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument(
        "--output-dir", type=Path, default=Path(__file__).parent / "direct-results"
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if args.warmups < 0 or args.rounds < 1 or args.max_tokens < 1:
        raise SystemExit("warmups must be >= 0; rounds and max-tokens must be >= 1")

    run_id = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
    output_dir = args.output_dir / f"{run_id}-{args.label}"
    output_dir.mkdir(parents=True, exist_ok=False)

    load_started = time.perf_counter()
    model, tokenizer = load(args.model, trust_remote_code=True)
    load_s = time.perf_counter() - load_started
    sampler = make_sampler(temp=0.0)
    prompt_items = list(PROMPTS.items())
    rows: list[dict[str, Any]] = []

    for phase, count in (("warmup", args.warmups), ("measured", args.rounds)):
        for round_number in range(count):
            prompt_name, prompt = prompt_items[round_number % len(prompt_items)]
            mx.random.seed(args.seed)
            started = time.perf_counter()
            first_token_at: float | None = None
            final = None
            error = None
            try:
                for response in stream_generate(
                    model,
                    tokenizer,
                    prompt,
                    max_tokens=args.max_tokens,
                    sampler=sampler,
                ):
                    if first_token_at is None:
                        first_token_at = time.perf_counter()
                    final = response
            except Exception as exception:  # Preserve evidence before failing the run.
                error = f"{type(exception).__name__}: {exception}"
            total_s = time.perf_counter() - started
            tokens = final.generation_tokens if final is not None else 0
            row = {
                "runtime": "mlx-lm-direct",
                "model": args.label,
                "source": args.model,
                "phase": phase,
                "round": round_number + 1,
                "prompt_name": prompt_name,
                "prompt_tokens": final.prompt_tokens if final is not None else None,
                "output_tokens": tokens,
                "ttft_s": first_token_at - started if first_token_at is not None else None,
                "total_s": total_s,
                "prompt_tps": final.prompt_tps if final is not None else None,
                "decode_tps": final.generation_tps if final is not None else None,
                "request_tps": tokens / total_s if tokens else None,
                "peak_memory_gb": final.peak_memory if final is not None else None,
                "finish_reason": final.finish_reason if final is not None else None,
                "error": error,
            }
            rows.append(row)
            print(
                f"{phase} {round_number + 1}/{count} prompt={prompt_name} "
                f"total={total_s:.3f}s tokens={tokens} "
                f"decode={row['decode_tps'] or 0:.2f}tok/s error={error or 'none'}",
                flush=True,
            )

    summary = summarize(rows)
    metadata = {
        "run_id": run_id,
        "created_at": datetime.now(UTC).isoformat(),
        "host": {
            "platform": platform.platform(),
            "machine": platform.machine(),
            "python": platform.python_version(),
        },
        "arguments": vars(args) | {"output_dir": str(args.output_dir)},
        "load_s": load_s,
        "mlx_version": getattr(mx, "__version__", None),
        "device": str(mx.default_device()),
        "prompts": PROMPTS,
        "summary": summary,
    }
    (output_dir / "raw.json").write_text(json.dumps(rows, indent=2), encoding="utf-8")
    (output_dir / "metadata.json").write_text(
        json.dumps(metadata, indent=2, default=str), encoding="utf-8"
    )
    (output_dir / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print(json.dumps(summary, indent=2))
    print(f"Results: {output_dir}")
    return 1 if any(row["error"] for row in rows if row["phase"] == "measured") else 0


if __name__ == "__main__":
    raise SystemExit(main())