#!/usr/bin/env python3
"""chronos-oracle-llm-tools · bench — measure an LLM's speed and response through Kairos moments, stamped by Chronos.

Each prompt is one action Kairos sequences (trigger → request → first_token → complete), timed on the
monotonic clock and stamped from a chronos.agent PromisedTime anchor (agents/kairos_meter.py). Tokens come
from the engine's own counts. Records go to data/logs/kairos_moments.jsonl; a summary prints at the end.

  python scripts/kairos_bench.py --engine openai --base http://127.0.0.1:8091 --model Bonsai-8B-Q1_0
  python scripts/kairos_bench.py --engine ollama --base http://127.0.0.1:11434 --model mindXtrain39
"""
from __future__ import annotations

import argparse
import asyncio
import json
import sys
import urllib.request
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from kairos_meter import KairosMeter, summarize  # noqa: E402

PROMPTS = [
    "What is 17 multiplied by 3? Answer with the number only.",
    "What is the capital of France? Answer in one word.",
    "In two sentences, explain why a 1-bit model uses less memory than a 16-bit model.",
]


def _stream(engine: str, base: str, model: str, prompt: str, max_tokens: int):
    """Yield ('token', text) for each piece and ('usage', {...}) at the end — blocking; run in a thread."""
    if engine == "openai":
        body = {"model": model, "messages": [{"role": "user", "content": prompt}], "max_tokens": max_tokens,
                "temperature": 0, "stream": True, "stream_options": {"include_usage": True}}
        req = urllib.request.Request(base.rstrip("/") + "/v1/chat/completions", data=json.dumps(body).encode(),
                                     headers={"Content-Type": "application/json"})
        usage = {}
        with urllib.request.urlopen(req, timeout=1800) as r:
            for raw in r:
                line = raw.decode("utf-8", "replace").strip()
                if not line.startswith("data:") or line == "data: [DONE]":
                    continue
                ev = json.loads(line[5:])
                for ch in ev.get("choices") or []:
                    t = (ch.get("delta") or {}).get("content")
                    if t:
                        yield "token", t
                usage = ev.get("usage") or usage
        yield "usage", {"prompt_tokens": usage.get("prompt_tokens"), "completion_tokens": usage.get("completion_tokens")}
    else:  # ollama native
        body = {"model": model, "messages": [{"role": "user", "content": prompt}], "stream": True,
                "options": {"num_predict": max_tokens, "temperature": 0}}
        req = urllib.request.Request(base.rstrip("/") + "/api/chat", data=json.dumps(body).encode(),
                                     headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=1800) as r:
            for raw in r:
                ev = json.loads(raw.decode("utf-8", "replace") or "{}")
                t = (ev.get("message") or {}).get("content")
                if t:
                    yield "token", t
                if ev.get("done"):
                    yield "usage", {"prompt_tokens": ev.get("prompt_eval_count"), "completion_tokens": ev.get("eval_count")}


async def run(a) -> int:
    meter = await KairosMeter.get_instance(engine_pid=a.engine_pid)
    loop = asyncio.get_running_loop()
    for i, prompt in enumerate(PROMPTS):
        q: asyncio.Queue = asyncio.Queue()

        def pump():
            try:
                for item in _stream(a.engine, a.base, a.model, prompt, a.max_tokens):
                    loop.call_soon_threadsafe(q.put_nowait, item)
            except Exception as e:  # noqa: BLE001
                loop.call_soon_threadsafe(q.put_nowait, ("error", e))
            loop.call_soon_threadsafe(q.put_nowait, ("end", None))

        text = []
        try:
            async with meter.moment("llm.chat", provider=a.engine, model=a.model, caller="kairos_bench",
                                    trigger=f"kairos_bench#{i}", sampling={"temperature": 0},
                                    model_digest=a.model_digest) as m:
                loop.run_in_executor(None, pump)
                while True:
                    kind, val = await q.get()
                    if kind == "token":
                        m.token(val)
                        text.append(val)
                    elif kind == "usage":
                        m.usage(**{k: v for k, v in val.items() if v is not None})
                    elif kind == "error":
                        raise val
                    elif kind == "end":
                        break
        except Exception as e:  # the moment already recorded the failure
            print(f"#{i} FAILED {e}")
            continue
        r = m.record
        print(f"#{r.seq} {a.model}: ttft {r.ttft_ms} ms · gen {r.gen_ms} ms · wall {r.wall_ms} ms · "
              f"{r.prompt_tokens}→{r.completion_tokens} tokens ({r.tokens_source}) · {r.completion_tok_per_s} tok/s · "
              f"prompt {r.prompt_tok_per_s} tok/s · cpu {r.cpu_s} s · gpu {r.gpu_ms} ms · start {r.start_unix_18dp} ({r.time_source}, {r.chronos_consensus}, "
              f"anchor age {r.anchor_age_s} s) | {''.join(text)[:60]!r}")
    print(json.dumps(summarize(last=len(PROMPTS)), indent=1))
    print("chronos anchors taken:", meter.anchor_refreshes)
    return 0


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--engine", choices=["openai", "ollama"], default="openai")
    ap.add_argument("--base", default="http://127.0.0.1:8091")
    ap.add_argument("--model", default="Bonsai-8B-Q1_0")
    ap.add_argument("--max-tokens", type=int, default=64)
    ap.add_argument("--engine-pid", type=int, default=None, help="the local engine's pid, for its CPU-seconds per response")
    ap.add_argument("--model-digest", default=None, help="the pinned model's sha256 (fate), e.g. from FORK.json")
    raise SystemExit(asyncio.run(run(ap.parse_args())))