#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""llama-swap + llama.cpp  ->  Prometheus-Text fuer Telegraf (inputs.exec).

Warum dieses Skript ueberhaupt:
  llama-swap vergibt die Upstream-Ports DYNAMISCH (`${PORT}`, startPort 9340,
  stabile Modell-ID-Reihenfolge) - ein statischer Scrape kann sie nicht
  treffen. Ausserdem liegen die Token-Zahlen in zwei Quellen: llama.cpp
  (/metrics, nur mit --metrics, kein Modell-Label) und llama-swaps eigenem
  Statistik-Store (/api/metrics/stats, das die Upstream-Antworten selbst
  auswertet und damit auch strata und sd-server abdeckt).
  Ein exec-Input, das llama-swap selbst fragt, ist deshalb der einzige Weg
  ohne Raten.

Quellen (alle nur lesend, jede mit eigenem Timeout, Fehler brechen nichts ab):
  * GET /running                        laufende Modelle + Upstream-Adresse
  * GET /v1/models                      alle Modell-IDs
  * GET /api/metrics/stats?model=<id>   Token-Summen je Modell
  * GET http://<upstream>/metrics       llama.cpp-Zaehler je Profil

Ausgabe: Prometheus-Textformat auf stdout.
  * Keine erreichbare Quelle = keine Zeile fuer diese Quelle.
  * Es wird nie ein Wert erfunden und nie ein Teil-Fehler als 0 gemeldet.
  * Rueckgabewert ist immer 0, damit Telegraf den Lauf nicht als Fehler sieht.

Aufruf durch Telegraf (Nutzer-Instanz auf dem ROG):
  [[inputs.exec]]
    commands = ["/home/mm/bin/llamaswap-telemetrie.py"]
    data_format = "prometheus"
    timeout = "25s"
    interval = "60s"
"""

import json
import os
import re
import sys
import urllib.error
import urllib.parse
import urllib.request

BASE = os.environ.get("LLAMASWAP_BASE", "http://192.168.188.184:9300")
TIMEOUT = float(os.environ.get("LLAMASWAP_TIMEOUT", "4"))

# Profile, die model-manager.sh auf festen Ports startet (Dauerlaeufer und
# Handstarts). Sie erscheinen NICHT in /running, wenn llama-swap sie nicht
# verwaltet - deshalb fest hinterlegt. Nicht erreichbare Ports kosten nur einen
# sofortigen ECONNREFUSED, keine Wartezeit.
FESTE_PORTS = {
    8101: "chat",
    8102: "embedding",
    8103: "vision",
    8104: "ocr",
    8105: "rerank",
    8107: "fast",
    8108: "agent",
    8109: "desktop",
}

METRIK_RE = re.compile(r"^([a-zA-Z_:][a-zA-Z0-9_:]*)(\{[^}]*\})?\s+(\S+)\s*$")


def hole(pfad, timeout=None):
    """GET gegen llama-swap. Gibt (status, bytes) zurueck - status 0 bei Fehler."""
    url = pfad if pfad.startswith("http") else BASE + pfad
    try:
        with urllib.request.urlopen(url, timeout=timeout or TIMEOUT) as antwort:
            return antwort.status, antwort.read()
    except urllib.error.HTTPError as e:
        return e.code, b""
    except Exception:
        return 0, b""


def hole_json(pfad):
    status, rohdaten = hole(pfad)
    if status != 200 or not rohdaten:
        return None
    try:
        return json.loads(rohdaten)
    except ValueError:
        return None


def name_saeubern(name):
    """llama.cpp liefert `llamacpp:prompt_tokens_total` (Doppelpunkt erlaubt,
    aber Prometheus reserviert ihn fuer Recording-Regeln) -> Unterstrich."""
    return name.replace(":", "_")


def doppelpunkt_label_frei(labels):
    """Label-String normalisieren: fuehrende/abschliessende Klammern entfernen."""
    if not labels:
        return ""
    labels = labels.strip()
    if labels.startswith("{") and labels.endswith("}"):
        return labels[1:-1]
    return ""


def upstream_kandidaten():
    """{port: modell} aus /running plus die festen model-manager-Ports."""
    kandidaten = dict(FESTE_PORTS)
    running = hole_json("/running")
    if isinstance(running, dict):
        for eintrag in running.get("running", []) or []:
            modell = eintrag.get("model") or ""
            proxy = eintrag.get("proxy") or ""
            m = re.search(r":(\d+)\s*$", proxy)
            if m and modell:
                # /running gewinnt: llama-swap kennt den echten Modellnamen.
                kandidaten[int(m.group(1))] = modell
    return kandidaten


def llama_cpp_zeilen(port, modell):
    """Prometheus-Zeilen eines llama.cpp-Upstreams, mit model-Label ergaenzt."""
    status, rohdaten = hole("http://192.168.188.184:%d/metrics" % port, timeout=2)
    if status != 200 or not rohdaten:
        return [], []
    text = rohdaten.decode("utf-8", "replace")
    if "llamacpp" not in text:          # kein llama.cpp (strata, sd-server)
        return [], []

    typen, zeilen = {}, []
    for zeile in text.splitlines():
        zeile = zeile.strip()
        if not zeile:
            continue
        if zeile.startswith("# TYPE"):
            teile = zeile.split()
            if len(teile) >= 4:
                typen[name_saeubern(teile[2])] = teile[3]
            continue
        if zeile.startswith("#"):
            continue
        m = METRIK_RE.match(zeile)
        if not m:
            continue
        name = name_saeubern(m.group(1))
        labels = doppelpunkt_label_frei(m.group(2))
        wert = m.group(3)
        if wert in ("NaN", "+Inf", "-Inf"):
            continue
        try:
            float(wert)
        except ValueError:
            continue
        neue_labels = ('model="%s"' % modell) + ("," + labels if labels else "")
        zeilen.append((name, neue_labels, wert))
    return typen, zeilen


def aktivitaets_zeilen():
    """Token-Summen je Modell aus llama-swaps Statistik-Store."""
    modelle = hole_json("/v1/models")
    if not isinstance(modelle, dict):
        return [], []
    ids = [m.get("id") for m in modelle.get("data", []) or [] if m.get("id")]
    felder = [
        ("total_requests", "llamaswap_activity_requests_total"),
        ("total_input_tokens", "llamaswap_activity_input_tokens_total"),
        ("total_output_tokens", "llamaswap_activity_output_tokens_total"),
        ("total_cache_tokens", "llamaswap_activity_cache_tokens_total"),
    ]
    zeilen = []
    for modell in ids:
        stats = hole_json("/api/metrics/stats?model=%s" % urllib.parse.quote(modell))
        if not isinstance(stats, dict):
            continue
        for schluessel, metrik in felder:
            wert = stats.get(schluessel)
            if isinstance(wert, (int, float)) and not isinstance(wert, bool):
                zeilen.append((metrik, 'model="%s"' % modell, "%d" % wert))
    return zeilen


def main():
    typen = {}
    zeilen = []

    kandidaten = upstream_kandidaten()
    for port, modell in sorted(kandidaten.items()):
        t, z = llama_cpp_zeilen(port, modell)
        typen.update(t)
        zeilen.extend(z)

    aktivitaet = aktivitaets_zeilen()
    for metrik, _, _ in aktivitaet:
        typen[metrik] = "counter"
    zeilen.extend(aktivitaet)

    if not zeilen:
        # Nichts erreichbar: lieber gar nichts ausgeben als Nullen.
        return 0

    ausgabe = []
    for name in sorted({n for n, _, _ in zeilen}):
        typ = typen.get(name, "gauge")
        ausgabe.append("# TYPE %s %s" % (name, typ))
        for n, labels, wert in zeilen:
            if n != name:
                continue
            ausgabe.append("%s{%s} %s" % (n, labels, wert))
    sys.stdout.write("\n".join(ausgabe) + "\n")
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except Exception as fehler:              # fail-open: nie den exec brechen
        sys.stderr.write("llamaswap-telemetrie: %s\n" % fehler)
        sys.exit(0)