#!/usr/bin/env python3
"""
Parse llama-server server log and extract timing data.

Usage:
    python3 parse-benchmark-log.py <logfile> [model_name]

The logfile is the stdout/stderr output produced by llama-server
(--metrics flag).  The script finds every prompt-eval and
generation-eval line and reports tokens/second for each run.

Requires no dependencies beyond the Python standard library.
"""

import re
import sys


def main() -> None:
    if len(sys.argv) < 2:
        print(__doc__.strip())
        sys.exit(1)

    logfile = sys.argv[1]
    name = sys.argv[2] if len(sys.argv) > 2 else "unknown"

    with open(logfile) as f:
        lines = f.readlines()

    prompt_results: list[tuple[str, str]] = []
    gen_results: list[tuple[str, str]] = []

    PROMPT_RE = re.compile(
        r"prompt eval time =\s+\S+\s+ms\s+/\s+(\d+)\s+tokens\s+"
        r"\(\s+\S+\s+ms per token,\s+([\d.]+)\s+tokens per second\)"
    )
    GEN_RE = re.compile(
        r"eval time =\s+\S+\s+ms\s+/\s+(\d+)\s+tokens\s+"
        r"\(\s+\S+\s+ms per token,\s+([\d.]+)\s+tokens per second\)"
    )

    for line in lines:
        pm = PROMPT_RE.search(line)
        if pm:
            prompt_results.append((pm.group(1), pm.group(2)))

        # Skip prompt-eval lines so GEN_RE does not double-match them
        # (the string "eval time =" appears inside "prompt eval time =")
        if line.strip().startswith("prompt"):
            continue
        gm = GEN_RE.search(line)
        if gm:
            gen_results.append((gm.group(1), gm.group(2)))

    print(f"Results for: {name}")
    print()

    print("Prompt eval runs:")
    for tok, tps in prompt_results:
        print(f"  prompt={tok} tok, {tps} tok/s")

    print()
    print("Generation runs:")
    for tok, tps in gen_results:
        print(f"  gen={tok} tok @ {tps} tok/s")

    print()
    try:
        pe = f"{prompt_results[0][1]}/{prompt_results[0][0]}"
    except IndexError:
        pe = "?/?"
    lines = len(gen_results)
    g1 = f"{gen_results[1][1]}/{gen_results[1][0]}" if lines > 1 else "?/?"
    g2 = f"{gen_results[2][1]}/{gen_results[2][0]}" if lines > 2 else "?/?"
    print(f"Summary:  prompt={pe}  |  gen300={g2}  |  gen150={g1}")


if __name__ == "__main__":
    main()