#!/usr/bin/env python3
"""Measure end-to-end latency for deterministic, complex programming tasks."""

from __future__ import annotations

import argparse
import hashlib
import json
import platform
import subprocess
import sys
import time
import urllib.error
import urllib.request
from dataclasses import asdict, dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

TASKS = {
    "python_async_cache": """
Act as a senior Python engineer. Implement a production-ready, single-file Python 3.12 module
containing a generic AsyncTTLCache[K, V]. Requirements: weighted LRU capacity; per-entry TTL based
on time.monotonic; single-flight get_or_create so concurrent misses for one key invoke the async
factory exactly once; cancellation of one waiter must not cancel the shared factory; exceptions are
never cached; invalidating a key while its factory is running must prevent stale publication; never
hold the cache lock while awaiting user code; clear(), invalidate(), and immutable hit/miss/eviction
statistics. Include concise pytest-asyncio tests covering races, cancellation, expiry, weighted
capacity, factory failure, and invalidation during creation. Use only the standard library plus
pytest. Output only the complete module, without Markdown fences or prose.
""".strip(),
    "php_idempotent_transfer": """
Act as a senior Laravel 12 engineer. Produce one self-contained PHP file that demonstrates an
idempotent money-transfer service and its PHPUnit tests. Store money as integer cents. Requirements:
validate positive amount and distinct accounts; use a database transaction; lock both account rows
in deterministic ID order; reject insufficient funds; use a unique client idempotency key so retries
return the original transfer; retry deadlocks without duplicating effects; persist an outbox event
in the same transaction; and prove with tests that duplicate and competing requests preserve
balances, create one transfer, and create one outbox event. Include minimal model/table assumptions
as comments.
Output only valid PHP, without Markdown fences or prose.
""".strip(),
    "typescript_dag_scheduler": """
Act as a senior TypeScript engineer. Implement a strict-mode, dependency-free DAG task scheduler in
one TypeScript file. Each task has an ID, dependency IDs, and an async run(signal) function.
Requirements: validate unknown dependencies and duplicate IDs; report an explicit cycle path; stable
input-order scheduling; configurable maximum concurrency; AbortSignal support; fail fast after the
first error while awaiting already-running tasks; never start dependents of failed tasks; return a
Map of typed results on success; and avoid unhandled promise rejections. Include a compact test
harness covering a diamond graph, concurrency enforcement, stable order, cycle detection, external
cancellation, and rejection. Output only TypeScript, without Markdown fences or prose.
""".strip(),
}


@dataclass(frozen=True)
class Runtime:
    name: str
    kind: str
    base_url: str
    model: str


@dataclass
class Result:
    task: str
    runtime: str
    model: str
    started_at: str
    prompt_chars: int
    max_tokens: int
    ttft_s: float | None = None
    total_s: float | None = None
    prompt_tokens: int | None = None
    output_tokens: int | None = None
    decode_tps: float | None = None
    finish_reason: str | None = None
    output_chars: int = 0
    output_sha256: str | None = None
    output: str = ""
    error: str | None = None


def utc_now() -> str:
    return datetime.now(UTC).isoformat()


def post_json(url: str, payload: dict[str, Any], timeout: float) -> dict[str, Any]:
    request = urllib.request.Request(
        url,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=timeout) as response:
        return json.load(response)


def health(runtime: Runtime) -> None:
    path = "/api/tags" if runtime.kind == "ollama" else "/v1/models"
    with urllib.request.urlopen(f"{runtime.base_url.rstrip('/')}{path}", timeout=15):
        pass


def consume_ollama(response: Any, started: float) -> tuple[float | None, str, dict[str, Any]]:
    first_at: float | None = None
    chunks: list[str] = []
    terminal: dict[str, Any] = {}
    for line in response:
        if not line.strip():
            continue
        event = json.loads(line)
        text = event.get("message", {}).get("content") or event.get("response") or ""
        if text:
            if first_at is None:
                first_at = time.perf_counter()
            chunks.append(text)
        if event.get("done"):
            terminal = event
    return (first_at - started if first_at else None, "".join(chunks), terminal)


def consume_sse(
    response: Any, started: float
) -> tuple[float | None, str, dict[str, Any], str | None]:
    first_at: float | None = None
    chunks: list[str] = []
    terminal: dict[str, Any] = {}
    finish_reason: str | None = None
    for raw_line in response:
        line = raw_line.decode("utf-8").strip()
        if not line.startswith("data:"):
            continue
        payload = line[5:].strip()
        if payload == "[DONE]":
            break
        event = json.loads(payload)
        terminal = event
        for choice in event.get("choices") or []:
            text = choice.get("text") or choice.get("delta", {}).get("content") or ""
            if text:
                if first_at is None:
                    first_at = time.perf_counter()
                chunks.append(text)
            finish_reason = choice.get("finish_reason") or finish_reason
    return (first_at - started if first_at else None, "".join(chunks), terminal, finish_reason)


def request_payload(
    runtime: Runtime, prompt: str, max_tokens: int, seed: int
) -> tuple[str, dict[str, Any]]:
    if runtime.kind == "ollama":
        return f"{runtime.base_url.rstrip('/')}/api/chat", {
            "model": runtime.model,
            "messages": [
                {
                    "role": "system",
                    "content": (
                        "You are a senior software engineer. Follow every output constraint "
                        "exactly and return a complete solution."
                    ),
                },
                {"role": "user", "content": prompt},
            ],
            "think": False,
            "stream": True,
            "keep_alive": "30m",
            "options": {
                "temperature": 0,
                "seed": seed,
                "num_predict": max_tokens,
                "num_ctx": 8192,
                "draft_num_predict": 0,
            },
        }
    return f"{runtime.base_url.rstrip('/')}/v1/chat/completions", {
        "model": runtime.model,
        "messages": [
            {
                "role": "system",
                "content": (
                    "You are a senior software engineer. Follow every output constraint exactly "
                    "and return a complete solution."
                ),
            },
            {"role": "user", "content": prompt},
        ],
        "stream": True,
        "stream_options": {"include_usage": True},
        "chat_template_kwargs": {"enable_thinking": False},
        "temperature": 0,
        "seed": seed,
        "max_tokens": max_tokens,
    }


def run_task(
    runtime: Runtime,
    task_name: str,
    prompt: str,
    max_tokens: int,
    seed: int,
    timeout: float,
) -> Result:
    result = Result(
        task=task_name,
        runtime=runtime.name,
        model=runtime.model,
        started_at=utc_now(),
        prompt_chars=len(prompt),
        max_tokens=max_tokens,
    )
    url, payload = request_payload(runtime, prompt, max_tokens, seed)
    request = urllib.request.Request(
        url,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json", "Accept": "text/event-stream"},
        method="POST",
    )
    started = time.perf_counter()
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            if runtime.kind == "ollama":
                result.ttft_s, result.output, terminal = consume_ollama(response, started)
                result.prompt_tokens = terminal.get("prompt_eval_count")
                result.output_tokens = terminal.get("eval_count")
                result.finish_reason = terminal.get("done_reason")
                decode_ns = terminal.get("eval_duration")
                if decode_ns and result.output_tokens:
                    result.decode_tps = result.output_tokens / (decode_ns / 1e9)
            else:
                result.ttft_s, result.output, terminal, result.finish_reason = consume_sse(
                    response, started
                )
                usage = terminal.get("usage") or {}
                result.prompt_tokens = usage.get("prompt_tokens")
                result.output_tokens = usage.get("completion_tokens")
            result.total_s = time.perf_counter() - started
            if result.decode_tps is None and result.output_tokens and result.ttft_s is not None:
                decode_s = result.total_s - result.ttft_s
                if decode_s > 0:
                    result.decode_tps = result.output_tokens / decode_s
            result.output_chars = len(result.output)
            result.output_sha256 = hashlib.sha256(result.output.encode()).hexdigest()
            if not result.output:
                result.error = "empty generated output"
    except urllib.error.HTTPError as error:
        detail = error.read().decode("utf-8", errors="replace")[:2000]
        result.error = f"HTTP {error.code}: {detail}"
        result.total_s = time.perf_counter() - started
    except (OSError, ValueError, json.JSONDecodeError) as error:
        result.error = f"{type(error).__name__}: {error}"
        result.total_s = time.perf_counter() - started
    return result


def package_versions() -> dict[str, str | None]:
    script = (
        "import importlib.metadata as m,json; "
        "print(json.dumps({p:m.version(p) for p in "
        "['vllm','vllm-metal','mlx','mlx-lm','mlx-vlm']}))"
    )
    executable = Path.home() / ".venv-vllm-metal/bin/python"
    try:
        process = subprocess.run(
            [str(executable), "-c", script], capture_output=True, text=True, timeout=30, check=True
        )
        return json.loads(process.stdout)
    except (OSError, subprocess.SubprocessError, json.JSONDecodeError):
        return {}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--runtime", choices=("ollama", "vllm"), required=True)
    parser.add_argument("--ollama-url", default="http://127.0.0.1:11434")
    parser.add_argument("--ollama-model", default="qwen3.8:27b-mlx")
    parser.add_argument("--vllm-url", default="http://127.0.0.1:8001")
    parser.add_argument("--vllm-model", default="qwen3.8-27b-nvfp4")
    parser.add_argument("--max-tokens", type=int, default=1536)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--timeout", type=float, default=900)
    parser.add_argument("--output-dir", type=Path, default=Path(__file__).parent / "coding-results")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if args.max_tokens < 1:
        raise SystemExit("max-tokens must be positive")
    runtime = (
        Runtime("ollama", "ollama", args.ollama_url, args.ollama_model)
        if args.runtime == "ollama"
        else Runtime("vllm-metal", "openai", args.vllm_url, args.vllm_model)
    )
    try:
        health(runtime)
    except Exception as error:
        raise SystemExit(f"{runtime.name} is unavailable: {error}") from error

    warmup = run_task(runtime, "warmup", "Output only the word READY.", 8, args.seed, args.timeout)
    if warmup.error:
        raise SystemExit(f"warmup failed: {warmup.error}")
    print(
        f"warmup runtime={runtime.name} total={warmup.total_s:.3f}s "
        f"tokens={warmup.output_tokens}",
        flush=True,
    )

    results: list[Result] = []
    for index, (task_name, prompt) in enumerate(TASKS.items(), start=1):
        result = run_task(runtime, task_name, prompt, args.max_tokens, args.seed, args.timeout)
        results.append(result)
        print(
            f"[{index}/{len(TASKS)}] task={task_name} total={result.total_s:.3f}s "
            f"ttft={result.ttft_s:.3f}s tokens={result.output_tokens} "
            f"decode={result.decode_tps:.2f}tok/s finish={result.finish_reason} "
            f"error={result.error or 'none'}",
            flush=True,
        )

    run_id = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
    output_dir = args.output_dir / f"{run_id}-{runtime.name}-{runtime.model.replace(':', '_')}"
    output_dir.mkdir(parents=True, exist_ok=False)
    raw = {
        "run_id": run_id,
        "created_at": utc_now(),
        "runtime": asdict(runtime),
        "host": {
            "platform": platform.platform(),
            "machine": platform.machine(),
            "python": platform.python_version(),
        },
        "package_versions": package_versions(),
        "settings": {
            "temperature": 0,
            "seed": args.seed,
            "max_tokens": args.max_tokens,
            "context_tokens": 8192,
            "raw_completion": False,
            "chat_template": True,
            "thinking": False,
            "ollama_speculative_drafting": False,
        },
        "warmup": asdict(warmup),
        "tasks": TASKS,
        "results": [asdict(result) for result in results],
    }
    (output_dir / "results.json").write_text(json.dumps(raw, indent=2), encoding="utf-8")
    summary = [
        {
            key: value
            for key, value in asdict(result).items()
            if key not in {"output"}
        }
        for result in results
    ]
    (output_dir / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print(f"Results: {output_dir}")
    return 1 if any(result.error for result in results) else 0


if __name__ == "__main__":
    sys.exit(main())