#!/usr/bin/env python3
"""XeCores benchmark runner for OpenAI-compatible GPU inference servers.

Standard library only. Supports vLLM, SGLang, and llama.cpp servers through
/v1/completions. Local API credentials and endpoint addresses are never placed
in the submission payload.
"""
from __future__ import annotations

import argparse
import concurrent.futures
import dataclasses
import datetime as dt
import hashlib
import ipaddress
import json
import math
import os
import pathlib
import platform
import re
import statistics
import subprocess
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
from typing import Callable, Iterable, Optional

RUNNER_VERSION = "0.1.0"
SCHEMA_VERSION = "1.0"
SUPPORTED_CONCURRENCY = (1, 2, 4, 8, 16, 32)
DEFAULT_API = "https://xecores.com/api/v1"
KNOWN_VRAM = {"b50": 16.0, "b60": 24.0, "b65": 32.0, "b70": 32.0, "b580": 12.0}


@dataclasses.dataclass
class Sample:
    index: int
    prompt_tokens: int
    output_tokens: int
    ttft_ms: Optional[float]
    e2e_ms: Optional[float]
    token_timestamps_ms: list[float]
    error: Optional[str]

    @property
    def inter_token_latency_ms(self) -> Optional[float]:
        if len(self.token_timestamps_ms) < 2:
            return None
        gaps = [b - a for a, b in zip(self.token_timestamps_ms, self.token_timestamps_ms[1:])]
        return statistics.median(gaps)

    @property
    def tpot_ms(self) -> Optional[float]:
        if self.e2e_ms is None or self.ttft_ms is None or self.output_tokens < 2:
            return None
        return max(0.0, (self.e2e_ms - self.ttft_ms) / (self.output_tokens - 1))


@dataclasses.dataclass
class RunConfig:
    engine: str
    endpoint: str
    api_key: Optional[str]
    model: str
    model_revision: str
    quantization: str
    dtype: str
    prompt_tokens: int
    output_tokens: int
    concurrency: int
    warmups: int
    requests: int
    timeout_s: float
    temperature: float
    seed: int
    credit: Optional[str]
    credit_url: Optional[str]
    gpu_model: Optional[str]
    gpu_count: Optional[int]
    power_limit_w: Optional[float]
    vram_gib_per_gpu: Optional[float] = None
    engine_version: Optional[str] = None
    ignore_eos: bool = False


def parse_concurrency(value: str) -> list[int]:
    parsed: list[int] = []
    for raw in value.split(","):
        item = raw.strip().lower()
        if not item:
            continue
        if item.startswith("c"):
            item = item[1:]
        try:
            number = int(item)
        except ValueError as exc:
            raise ValueError("Concurrency must use C1, C2, C4, C8, C16, C32") from exc
        if number not in SUPPORTED_CONCURRENCY:
            raise ValueError("Concurrency must use C1, C2, C4, C8, C16, C32")
        if number not in parsed:
            parsed.append(number)
    if not parsed:
        raise ValueError("Select at least one of C1, C2, C4, C8, C16, C32")
    return parsed


def parse_int_list(value: str, name: str, minimum: int, maximum: int) -> list[int]:
    values: list[int] = []
    for raw in value.split(","):
        try:
            number = int(raw.strip())
        except ValueError as exc:
            raise ValueError(f"{name} must be comma-separated integers") from exc
        if number < minimum or number > maximum:
            raise ValueError(f"{name} must be between {minimum} and {maximum}")
        if number not in values:
            values.append(number)
    return values


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


def metric_pair(values: list[float]) -> dict[str, Optional[float]]:
    return {"p50": percentile(values, 0.50), "p95": percentile(values, 0.95)}


def summarize(samples: list[Sample], wall_time_s: float) -> dict:
    successful = [sample for sample in samples if not sample.error]
    failed = [sample for sample in samples if sample.error]
    wall = max(wall_time_s, 0.000001)
    input_tokens = sum(sample.prompt_tokens for sample in successful)
    output_tokens = sum(sample.output_tokens for sample in successful)
    ttft = [sample.ttft_ms for sample in successful if sample.ttft_ms is not None]
    tpot = [sample.tpot_ms for sample in successful if sample.tpot_ms is not None]
    e2e = [sample.e2e_ms for sample in successful if sample.e2e_ms is not None]
    itl = [sample.inter_token_latency_ms for sample in successful if sample.inter_token_latency_ms is not None]
    return {
        "successful_requests": len(successful),
        "failed_requests": len(failed),
        "wall_time_s": round(wall_time_s, 4),
        "request_throughput_rps": round(len(successful) / wall, 6),
        "input_tokens_per_s": round(input_tokens / wall, 6),
        "output_tokens_per_s": round(output_tokens / wall, 6),
        "total_tokens_per_s": round((input_tokens + output_tokens) / wall, 6),
        "ttft_ms": metric_pair(ttft),
        "tpot_ms": metric_pair(tpot),
        "e2e_ms": metric_pair(e2e),
        "inter_token_latency_ms": metric_pair(itl),
    }


def parse_sse(lines: Iterable[bytes], started: float, clock: Callable[[], float] = time.monotonic) -> dict:
    text_parts: list[str] = []
    timestamps: list[float] = []
    prompt_tokens = 0
    output_tokens = 0
    for raw in lines:
        line = raw.decode("utf-8", errors="replace").strip()
        if not line.startswith("data:"):
            continue
        data = line[5:].strip()
        if data == "[DONE]":
            break
        try:
            event = json.loads(data)
        except json.JSONDecodeError:
            continue
        usage = event.get("usage") or {}
        prompt_tokens = int(usage.get("prompt_tokens") or prompt_tokens)
        output_tokens = int(usage.get("completion_tokens") or output_tokens)
        for choice in event.get("choices") or []:
            piece = choice.get("text")
            if piece is None:
                piece = (choice.get("delta") or {}).get("content")
            if piece:
                text_parts.append(str(piece))
                timestamps.append(max(0.0, (clock() - started) * 1000.0))
    if output_tokens <= 0:
        output_tokens = len(timestamps)
    return {
        "text": "".join(text_parts),
        "prompt_tokens": prompt_tokens,
        "output_tokens": output_tokens,
        "token_timestamps_ms": timestamps,
    }


def deterministic_prompt(target_tokens: int, seed: int, request_index: int) -> str:
    words = [
        "intel", "arc", "matrix", "memory", "kernel", "driver", "inference", "batch",
        "latency", "throughput", "context", "token", "compute", "cache", "power", "vector",
    ]
    prefix = f"XeCores benchmark seed {seed} request {request_index}. Continue with concise technical text. "
    count = max(1, target_tokens - 14)
    body = " ".join(f"{words[(index + request_index) % len(words)]}{(index * 17 + seed) % 997}" for index in range(count))
    return prefix + body


def safe_error(exc: BaseException) -> str:
    if isinstance(exc, urllib.error.HTTPError):
        return f"http_{exc.code}"
    if isinstance(exc, TimeoutError):
        return "timeout"
    name = type(exc).__name__.lower()
    return re.sub(r"[^a-z0-9_-]", "_", name)[:120] or "request_failed"


def endpoint_url(base: str) -> str:
    parsed = urllib.parse.urlparse(base)
    if parsed.scheme not in {"http", "https"} or not parsed.hostname:
        raise ValueError("Endpoint must be an http(s) URL")
    if parsed.username or parsed.password:
        raise ValueError("Put API credentials in --api-key or OPENAI_API_KEY, not in the URL")
    path = parsed.path.rstrip("/")
    if path.endswith("/v1"):
        path += "/completions"
    elif not path.endswith("/v1/completions"):
        path += "/v1/completions"
    return urllib.parse.urlunparse((parsed.scheme, parsed.netloc, path, "", "", ""))


def check_url_boundary(url: str, allow_local: bool) -> str:
    """Validate an outbound URL before any request leaves this process.

    Enforces http/https, no embedded credentials, and a resolved-IP boundary:
    loopback/private/link-local/reserved targets are refused unless
    allow_local is set (the benchmark endpoint is operator-supplied on the CLI
    and legitimately points at engines on this machine). DNS is resolved here
    so a hostname that rebinding-resolves to an internal address is caught."""
    import socket
    parsed = urllib.parse.urlparse(url)
    if parsed.scheme not in {"http", "https"} or not parsed.hostname:
        raise ValueError(f"refusing non-http(s) URL: {url!r}")
    if parsed.username or parsed.password:
        raise ValueError("credentials belong in headers, not the URL")
    infos = socket.getaddrinfo(parsed.hostname, parsed.port or (443 if parsed.scheme == "https" else 80),
                               type=socket.SOCK_STREAM)
    for info in infos:
        ip = info[4][0]
        if ipaddress.ip_address(ip).is_loopback or ipaddress.ip_address(ip).is_private \
                or ipaddress.ip_address(ip).is_link_local or ipaddress.ip_address(ip).is_reserved:
            if not allow_local:
                raise ValueError(f"refusing URL whose host resolves to a non-public address ({ip}): {url!r}")
    return url


def safe_fetch(url: str, *, data: Optional[bytes] = None, headers: Optional[dict] = None,
               method: str = "GET", timeout: float = 30.0, allow_local: bool = False):
    """The only egress point of this runner: boundary-check, then request."""
    check_url_boundary(url, allow_local)
    request = urllib.request.Request(url, data=data, headers=headers or {}, method=method)
    return urllib.request.urlopen(request, timeout=timeout)


def run_request(config: RunConfig, index: int) -> Sample:
    started = time.monotonic()
    body = {
        "model": config.model,
        "prompt": deterministic_prompt(config.prompt_tokens, config.seed, index),
        "max_tokens": config.output_tokens,
        "temperature": config.temperature,
        "seed": config.seed + index,
        "stream": True,
        "stream_options": {"include_usage": True},
    }
    if config.ignore_eos:
        body["ignore_eos"] = True
    headers = {"Content-Type": "application/json", "Accept": "text/event-stream", "User-Agent": f"xecores-runner/{RUNNER_VERSION}"}
    if config.api_key:
        headers["Authorization"] = f"Bearer {config.api_key}"
    url = endpoint_url(config.endpoint)
    try:
        with safe_fetch(url, data=json.dumps(body).encode(), headers=headers,
                        method="POST", timeout=config.timeout_s, allow_local=True) as response:
            parsed = parse_sse(response, started)
        ended = time.monotonic()
        stamps = parsed["token_timestamps_ms"]
        ttft = stamps[0] if stamps else None
        e2e = (ended - started) * 1000.0
        return Sample(index, parsed["prompt_tokens"], parsed["output_tokens"], ttft, e2e, stamps, None)
    except Exception as exc:  # request errors become bounded sample codes
        return Sample(index, 0, 0, None, None, [], safe_error(exc))


def run_profile(config: RunConfig) -> tuple[list[Sample], float]:
    for warmup in range(config.warmups):
        warm = dataclasses.replace(config, concurrency=1)
        sample = run_request(warm, -(warmup + 1))
        if sample.error:
            raise RuntimeError(f"Warmup failed: {sample.error}")
    started = time.monotonic()
    with concurrent.futures.ThreadPoolExecutor(max_workers=config.concurrency) as pool:
        samples = list(pool.map(lambda index: run_request(config, index), range(config.requests)))
    return samples, time.monotonic() - started


def command_output(command: list[str], timeout: float = 4.0) -> str:
    try:
        result = subprocess.run(command, capture_output=True, text=True, timeout=timeout, check=False)
        return result.stdout.strip()
    except (OSError, subprocess.SubprocessError):
        return ""


def detect_os() -> str:
    try:
        content = pathlib.Path("/etc/os-release").read_text(errors="ignore")
        match = re.search(r'^PRETTY_NAME="?([^"\n]+)', content, re.MULTILINE)
        if match:
            return match.group(1)[:120]
    except OSError:
        pass
    return platform.system() or "unknown"


def detect_cpu() -> str:
    try:
        content = pathlib.Path("/proc/cpuinfo").read_text(errors="ignore")
        match = re.search(r"^(?:model name|Hardware)\s*:\s*(.+)$", content, re.MULTILINE)
        if match:
            return match.group(1).strip()[:160]
    except OSError:
        pass
    return platform.processor() or "unknown"


def detect_ram_gib() -> float:
    try:
        pages = os.sysconf("SC_PHYS_PAGES")
        page_size = os.sysconf("SC_PAGE_SIZE")
        return round(pages * page_size / 1024**3, 2)
    except (ValueError, OSError):
        return 1.0


def detect_gpu_model() -> Optional[str]:
    output = command_output(["lspci"])
    for line in output.splitlines():
        if "intel" in line.lower() and any(term in line.lower() for term in ["vga", "display", "3d controller"]):
            return line.split(": ", 1)[-1].strip()[:100]
    return None


def known_vram(model: str) -> Optional[float]:
    lower = model.lower()
    for key, value in KNOWN_VRAM.items():
        if re.search(rf"\b{re.escape(key)}\b", lower):
            return value
    return None


def detect_driver() -> str:
    for path in ["/sys/module/xe/version", "/sys/module/i915/version"]:
        try:
            value = pathlib.Path(path).read_text().strip()
            if value:
                return value[:100]
        except OSError:
            pass
    output = command_output(["modinfo", "xe"])
    match = re.search(r"^version:\s*(.+)$", output, re.MULTILINE)
    return match.group(1).strip()[:100] if match else "unknown"


def profile_id(config: RunConfig) -> str:
    return f"p{config.prompt_tokens}-g{config.output_tokens}-c{config.concurrency}-v1"


def sample_json(sample: Sample) -> dict:
    return {
        "index": sample.index,
        "prompt_tokens": sample.prompt_tokens,
        "output_tokens": sample.output_tokens,
        "ttft_ms": round(sample.ttft_ms, 4) if sample.ttft_ms is not None else None,
        "e2e_ms": round(sample.e2e_ms, 4) if sample.e2e_ms is not None else None,
        "tpot_ms": round(sample.tpot_ms, 4) if sample.tpot_ms is not None else None,
        "inter_token_latency_ms": round(sample.inter_token_latency_ms, 4) if sample.inter_token_latency_ms is not None else None,
        "error": sample.error,
    }


def build_payload(config: RunConfig, samples: list[Sample], wall_time_s: float, challenge: Optional[dict] = None) -> dict:
    gpu_model = config.gpu_model or detect_gpu_model() or "Intel GPU (model not detected)"
    vram = config.vram_gib_per_gpu or known_vram(gpu_model) or 1.0
    now = dt.datetime.now(dt.timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
    payload = {
        "schema_version": SCHEMA_VERSION,
        "runner": {
            "name": "xecores-runner", "version": RUNNER_VERSION, "run_id": str(uuid.uuid4()),
            "started_at": now, "duration_s": round(wall_time_s, 4),
        },
        "engine": {"name": config.engine, "version": config.engine_version or "unknown", "api": "openai-completions"},
        "hardware": {
            "gpu_vendor": "Intel", "gpu_model": gpu_model, "gpu_count": config.gpu_count or 1,
            "vram_gib_per_gpu": vram, "driver": detect_driver(), "kernel": platform.release() or "unknown",
            "os": detect_os(), "cpu": detect_cpu(), "ram_gib": detect_ram_gib(), "power_limit_w": config.power_limit_w,
        },
        "model": {"id": config.model, "revision": config.model_revision, "quantization": config.quantization, "dtype": config.dtype},
        "profile": {
            "id": profile_id(config), "mode": "closed-loop", "concurrency": config.concurrency,
            "prompt_target_tokens": config.prompt_tokens, "output_target_tokens": config.output_tokens,
            "warmup_requests": config.warmups, "measured_requests": config.requests, "stream": True,
            "temperature": config.temperature, "seed": config.seed,
        },
        "aggregate": summarize(samples, wall_time_s),
        "samples": [sample_json(sample) for sample in samples],
        "credit": ({"display_name": config.credit, "url": config.credit_url} if config.credit else None),
    }
    if challenge:
        payload["challenge"] = {"id": challenge["id"], "token": challenge["token"]}
    return payload


def request_json(url: str, method: str = "GET", payload: Optional[dict] = None, timeout: float = 30.0) -> dict:
    data = json.dumps(payload).encode() if payload is not None else None
    headers = {"Accept": "application/json", "User-Agent": f"xecores-runner/{RUNNER_VERSION}"}
    if data is not None:
        headers["Content-Type"] = "application/json"
    with safe_fetch(url, data=data, headers=headers, method=method, timeout=timeout) as response:
        return json.loads(response.read().decode("utf-8"))


def submit_payload(api_base: str, payload: dict) -> dict:
    base = api_base.rstrip("/")
    challenge = request_json(f"{base}/challenge")
    payload = dict(payload)
    payload["challenge"] = {"id": challenge["id"], "token": challenge["token"]}
    return request_json(f"{base}/results", "POST", payload, timeout=60.0)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="xecores-runner", description="Run repeatable Intel GPU inference benchmarks and optionally submit them to XeCores.")
    parser.add_argument("run", nargs="?", help=argparse.SUPPRESS)
    parser.add_argument("--engine", required=True, choices=["vllm", "sglang", "llama.cpp"])
    parser.add_argument("--endpoint", required=True, help="Local OpenAI-compatible server base URL, such as http://127.0.0.1:8000/v1")
    parser.add_argument("--api-key", default=os.getenv("OPENAI_API_KEY"), help="Server key; defaults to OPENAI_API_KEY and is never submitted")
    parser.add_argument("--model", required=True, help="Exact served model identifier")
    parser.add_argument("--model-revision", default="unknown")
    parser.add_argument("--quantization", required=True)
    parser.add_argument("--dtype", required=True)
    parser.add_argument("--concurrency", default="c1,c2,c4,c8,c16,c32")
    parser.add_argument("--prompt-tokens", default="512", help="Comma-separated target lengths, up to 131072")
    parser.add_argument("--output-tokens", default="128", help="Comma-separated target lengths, up to 8192")
    parser.add_argument("--warmups", type=int, default=2, help="Warmup requests per profile")
    parser.add_argument("--requests-per-worker", type=int, default=3, help="Measured requests per concurrent worker")
    parser.add_argument("--timeout", type=float, default=300.0)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--ignore-eos", action="store_true", help="Ask compatible servers to generate the full output length")
    parser.add_argument("--gpu-model", help="Override detected Intel GPU model")
    parser.add_argument("--gpu-count", type=int, help="Override detected GPU count")
    parser.add_argument("--vram-gib", type=float, help="VRAM per GPU")
    parser.add_argument("--power-limit-w", type=float)
    parser.add_argument("--engine-version", default="unknown")
    parser.add_argument("--credit", help="Optional public display name")
    parser.add_argument("--credit-url", help="Optional public HTTPS profile URL")
    parser.add_argument("--submit", action="store_true", help="Submit each completed profile as pending review")
    parser.add_argument("--api-base", default=DEFAULT_API)
    parser.add_argument("--output", default="xecores-results.json")
    return parser


def validate_args(args: argparse.Namespace) -> None:
    endpoint_url(args.endpoint)
    if args.warmups < 0 or args.warmups > 128:
        raise ValueError("--warmups must be between 0 and 128")
    if args.requests_per_worker < 1 or args.requests_per_worker > 8:
        raise ValueError("--requests-per-worker must be between 1 and 8")
    if args.gpu_count is not None and not 1 <= args.gpu_count <= 16:
        raise ValueError("--gpu-count must be between 1 and 16")
    if args.credit_url and not args.credit:
        raise ValueError("--credit-url requires --credit")
    if args.credit_url and not args.credit_url.startswith("https://"):
        raise ValueError("--credit-url must use HTTPS")


def main(argv: Optional[list[str]] = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    try:
        validate_args(args)
        concurrencies = parse_concurrency(args.concurrency)
        prompts = parse_int_list(args.prompt_tokens, "--prompt-tokens", 1, 131072)
        outputs = parse_int_list(args.output_tokens, "--output-tokens", 1, 8192)
    except ValueError as exc:
        parser.error(str(exc))

    all_payloads: list[dict] = []
    total_profiles = len(concurrencies) * len(prompts) * len(outputs)
    profile_number = 0
    for prompt_tokens in prompts:
        for output_tokens in outputs:
            for concurrency in concurrencies:
                profile_number += 1
                measured = concurrency * args.requests_per_worker
                config = RunConfig(
                    engine=args.engine, endpoint=args.endpoint, api_key=args.api_key, model=args.model,
                    model_revision=args.model_revision, quantization=args.quantization, dtype=args.dtype,
                    prompt_tokens=prompt_tokens, output_tokens=output_tokens, concurrency=concurrency,
                    warmups=args.warmups, requests=measured, timeout_s=args.timeout,
                    temperature=args.temperature, seed=args.seed, credit=args.credit,
                    credit_url=args.credit_url, gpu_model=args.gpu_model, gpu_count=args.gpu_count,
                    power_limit_w=args.power_limit_w, vram_gib_per_gpu=args.vram_gib,
                    engine_version=args.engine_version, ignore_eos=args.ignore_eos,
                )
                print(f"[{profile_number}/{total_profiles}] {profile_id(config)}: {measured} measured requests", flush=True)
                samples, wall = run_profile(config)
                payload = build_payload(config, samples, wall)
                all_payloads.append(payload)
                aggregate = payload["aggregate"]
                print(f"  {aggregate['successful_requests']} ok · {aggregate['failed_requests']} failed · {aggregate['output_tokens_per_s']:.2f} output tok/s · TTFT p50 {aggregate['ttft_ms']['p50']} ms", flush=True)
                if args.submit:
                    try:
                        result = submit_payload(args.api_base, payload)
                        print(f"  submitted: {result.get('id', 'accepted')} · status {result.get('status', 'pending')}", flush=True)
                    except Exception as exc:
                        print(f"  submission failed: {safe_error(exc)}; local result kept", file=sys.stderr)

    output = pathlib.Path(args.output)
    output.write_text(json.dumps({"runner_version": RUNNER_VERSION, "results": all_payloads}, indent=2) + "\n")
    checksum = hashlib.sha256(output.read_bytes()).hexdigest()
    print(f"Saved {len(all_payloads)} profile(s) to {output} · sha256 {checksum}")
    return 0 if all(payload["aggregate"]["successful_requests"] for payload in all_payloads) else 2


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