"""Parse what `llama-bench` prints, markdown table or `-o json`, into benchmark rows.

People paste this output, so it is untrusted text. Only the measured numbers and a few
settings are kept. The JSON form carries the model's full file path, CPU model and
build machine details; none of that is stored.
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass

MAX_INPUT = 40_000

_GPU_NAME = re.compile(
    r"(NVIDIA (?:GeForce )?(?:RTX|GTX|RTX PRO|Tesla|A|H|L)\s?[A-Za-z0-9 ]{1,24}?)(?=[,;|\n]|$| \()"
    r"|(AMD Radeon[A-Za-z0-9 ]{1,30}?)(?=[,;|\n]|$| \()"
    r"|(Apple M[1-9](?: (?:Pro|Max|Ultra))?)"
    r"|(Intel\(R\) Arc\(TM\) [A-Za-z0-9 ]{1,12}|Intel Arc [A-Za-z0-9 ]{1,12})",
)
_BACKENDS = ("CUDA", "ROCm", "Metal", "Vulkan", "SYCL", "CPU", "BLAS", "RPC", "MUSA", "CANN", "OpenCL")


@dataclass
class BenchRow:
    test: str                 # "pp512", "tg128", "pp512+tg128"
    n_prompt: int
    n_gen: int
    tps: float
    stdev: float | None
    backend: str | None
    n_gpu_layers: int | None
    model_type: str | None    # "qwen3moe 30B.A3B Q4_K - Medium"
    size_gib: float | None
    params_b: float | None
    n_threads: int | None
    flash_attn: bool | None
    type_k: str | None
    gpu_name: str | None
    build: str | None

    def to_dict(self) -> dict:
        return dict(self.__dict__)


def _backend(text: str | None) -> str | None:
    if not text:
        return None
    found = [b for b in _BACKENDS if b.lower() in text.lower()]
    found = [b for b in found if b not in ("BLAS", "RPC") or len(found) == 1]
    return ",".join(found[:2]) or None


def gpu_name(text: str | None) -> str | None:
    if not text:
        return None
    m = _GPU_NAME.search(text)
    if not m:
        return None
    name = next(g for g in m.groups() if g)
    return re.sub(r"\s+", " ", name).strip()[:40]


def _int(v):
    try:
        return int(v)
    except (TypeError, ValueError):
        return None


def _tests(n_prompt: int, n_gen: int) -> str:
    if n_prompt and n_gen:
        return f"pp{n_prompt}+tg{n_gen}"
    return f"pp{n_prompt}" if n_prompt else f"tg{n_gen}"


def parse_json(text: str) -> list[BenchRow]:
    data = json.loads(text)
    if isinstance(data, dict):
        data = [data]
    rows = []
    for d in data:
        if not isinstance(d, dict) or "avg_ts" not in d:
            continue
        n_prompt, n_gen = _int(d.get("n_prompt")) or 0, _int(d.get("n_gen")) or 0
        size = d.get("model_size")
        params = d.get("model_n_params")
        rows.append(BenchRow(
            test=_tests(n_prompt, n_gen), n_prompt=n_prompt, n_gen=n_gen, tps=float(d["avg_ts"]),
            stdev=float(d["stdev_ts"]) if d.get("stdev_ts") is not None else None,
            backend=_backend(str(d.get("backends") or "")), n_gpu_layers=_int(d.get("n_gpu_layers")),
            model_type=str(d.get("model_type") or "")[:60] or None,
            size_gib=round(size / 2**30, 2) if isinstance(size, (int, float)) else None,
            params_b=round(params / 1e9, 2) if isinstance(params, (int, float)) else None,
            n_threads=_int(d.get("n_threads")),
            flash_attn=bool(d["flash_attn"]) if d.get("flash_attn") is not None else None,
            type_k=str(d.get("type_k") or "")[:8] or None,
            gpu_name=gpu_name(str(d.get("gpu_info") or "")),
            build=str(d.get("build_commit") or "")[:12] or None,
        ))
    return rows


_ROW_SPLIT = re.compile(r"\s*\|\s*")


def parse_markdown(text: str) -> list[BenchRow]:
    lines = [l.strip() for l in text.splitlines()]
    header = None
    rows = []
    build = None
    m = re.search(r"build:\s*([0-9a-f]{6,12})", text)
    if m:
        build = m.group(1)
    gpu = gpu_name(text)
    for line in lines:
        if not line.startswith("|"):
            continue
        cells = [c.strip() for c in line.strip().strip("|").split("|")]
        if header is None:
            if "t/s" in cells and "test" in cells:
                header = [c.strip().lower() for c in cells]
            continue
        if set(line.replace("|", "").strip()) <= set("-: "):
            continue
        if len(cells) != len(header):
            continue
        rec = dict(zip(header, (c.strip() for c in cells)))
        test = rec.get("test", "")
        tm = re.fullmatch(r"(?:pp(\d+))?\+?(?:tg(\d+))?(?:\s*@\s*d\d+)?", test)
        ts = re.match(r"([0-9.]+)(?:\s*±\s*([0-9.]+))?", rec.get("t/s", ""))
        if not tm or not ts or not test:
            continue
        size = re.match(r"([0-9.]+)\s*GiB", rec.get("size", ""))
        params = re.match(r"([0-9.]+)\s*B", rec.get("params", ""))
        fa = rec.get("fa")
        rows.append(BenchRow(
            test=test.split("@")[0].strip(), n_prompt=_int(tm.group(1)) or 0, n_gen=_int(tm.group(2)) or 0,
            tps=float(ts.group(1)), stdev=float(ts.group(2)) if ts.group(2) else None,
            backend=_backend(rec.get("backend")), n_gpu_layers=_int(rec.get("ngl")),
            model_type=(rec.get("model") or "")[:60] or None,
            size_gib=float(size.group(1)) if size else None, params_b=float(params.group(1)) if params else None,
            n_threads=_int(rec.get("threads")), flash_attn=(fa == "1") if fa in ("0", "1") else None,
            type_k=(rec.get("type_k") or "")[:8] or None, gpu_name=gpu, build=build,
        ))
    return rows


def parse(text: str) -> list[BenchRow]:
    text = (text or "")[:MAX_INPUT]
    stripped = text.strip()
    if stripped.startswith(("[", "{")):
        try:
            return parse_json(stripped)
        except (ValueError, TypeError):
            pass
    return parse_markdown(text)


def summarize(rows: list[BenchRow]) -> dict | None:
    """The numbers the dataset stores: prompt processing and generation speed, one row."""
    if not rows:
        return None
    pp = next((r for r in rows if r.n_prompt and not r.n_gen), None)
    tg = next((r for r in rows if r.n_gen and not r.n_prompt), None)
    ref = tg or pp or rows[0]
    return {
        "prompt_tokens": pp.n_prompt if pp else None, "prompt_tps": pp.tps if pp else None,
        "generation_tokens": tg.n_gen if tg else None, "generation_tps": tg.tps if tg else None,
        "backend": ref.backend, "n_gpu_layers": ref.n_gpu_layers, "bench_model_type": ref.model_type,
        "bench_size_gib": ref.size_gib, "bench_params_b": ref.params_b, "n_threads": ref.n_threads,
        "flash_attn": ref.flash_attn, "type_k": ref.type_k, "gpu_name": ref.gpu_name, "build": ref.build,
    }