import argparse
import os
from time import sleep
from urllib.error import HTTPError
from urllib.request import ProxyHandler, Request, build_opener

from compatible_chat import NoRedirects, local_base

METRICS = {
    "llamacpp:requests_processing": "active",
    "llamacpp:requests_deferred": "waiting",
    "llamacpp:tokens_predicted_total": "generated_total",
}


def read_metrics(base):
    key = os.environ.get("LOCAL_MODEL_API_KEY")
    headers = {"Authorization": f"Bearer {key}"} if key else {}
    endpoint = base.removesuffix("/v1") + "/metrics"
    opener = build_opener(ProxyHandler({}), NoRedirects())
    with opener.open(Request(endpoint, headers=headers), timeout=5) as response:
        body = response.read().decode("utf-8")
    values = {}
    for line in body.splitlines():
        parts = line.split()
        if len(parts) >= 2 and parts[0] in METRICS:
            values[METRICS[parts[0]]] = float(parts[1])
    if len(values) != len(METRICS):
        raise ValueError("Expected llama.cpp metrics are missing. Check your server build.")
    return values


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Sample metrics from your own local llama.cpp server.")
    parser.add_argument("--base", type=local_base, default="http://127.0.0.1:8080/v1")
    parser.add_argument("--samples", type=int, choices=range(1, 31), default=12)
    args = parser.parse_args()
    try:
        for index in range(args.samples):
            values = read_metrics(args.base)
            print(f"{index + 1:02d} | active={values['active']:.0f} | waiting={values['waiting']:.0f} "
                  f"| generated_total={values['generated_total']:.0f}", flush=True)
            if index + 1 < args.samples:
                sleep(1)
    except HTTPError as error:
        raise SystemExit(f"Metrics returned HTTP {error.code}. Check --metrics and the local API key.") from None
    except (OSError, ValueError) as error:
        raise SystemExit(str(error)) from None