#!/usr/bin/env python3
"""LLM inference metrics exporter for Prometheus.

Background thread scrapes only loaded models to avoid triggering model loads
in llama-server, then the HTTP handler serves the cached data immediately.

Exposes metrics for every known model from /models; unloaded models show
zero values without ever requesting /metrics?model=<unloaded>.

Uses persistent JSON cache on disk to survive restarts and compute counter
deltas for Prometheus rate/irate queries.
"""

import json
import os
import time
import threading
import logging
import urllib.request
import urllib.error
from http.server import HTTPServer, BaseHTTPRequestHandler

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

POLL_INTERVAL = int(os.environ.get("LLAMA_EXPORTER_INTERVAL", "15"))
SCRAPE_TIMEOUT = int(os.environ.get("LLAMA_EXPORTER_TIMEOUT", "5"))
CACHE_FILE = os.environ.get(
    "LLAMA_EXPORTER_CACHE",
    os.path.join(os.path.dirname(os.path.abspath(__file__)), "llama_exporter_cache.json"),
)

LLAMA_CPP_DEFAULTS = [
    {"name": "llama.cpp", "url": "http://localhost:11434"},
]

# All llama.cpp metrics we expose per model (counters + gauges).
ALL_METRICS = [
    "llamacpp:prompt_tokens_total",
    "llamacpp:prompt_seconds_total",
    "llamacpp:tokens_predicted_total",
    "llamacpp:tokens_predicted_seconds_total",
    "llamacpp:n_decode_total",
    "llamacpp:n_tokens_max",
    "llamacpp:prompt_tokens_seconds",
    "llamacpp:predicted_tokens_seconds",
    "llamacpp:requests_processing",
    "llamacpp:requests_deferred",
    "llamacpp:n_busy_slots_per_decode",
]

# Counter metrics that accumulate over time — exposed as _delta for rate() queries.
COUNTER_METRICS = {
    "llamacpp:prompt_tokens_total",
    "llamacpp:prompt_seconds_total",
    "llamacpp:tokens_predicted_total",
    "llamacpp:tokens_predicted_seconds_total",
    "llamacpp:n_decode_total",
}


class Cache:
    """Persistent, thread-safe metrics cache with delta computation."""

    def __init__(self):
        self._lock = threading.Lock()
        self._data = {}  # { model_id: { metric_name: value } }
        self._known_models = set()
        self._health = None
        self._loaded = {}

        # Previous state for delta computation: { model_id: { metric_name: value } }
        self._prev = {}
        self._prev_timestamp = 0.0
        self._deltas = {}

        # Load persisted state
        self._load_cache()

    def _load_cache(self):
        """Load previous scrape state from disk."""
        try:
            with open(CACHE_FILE, "r") as f:
                state = json.load(f)
            self._prev = state.get("models", {})
            self._prev_timestamp = state.get("timestamp", 0.0)
            logger.info("Loaded cache from %s (timestamp=%s, models=%d)",
                        CACHE_FILE, self._prev_timestamp, len(self._prev))
        except FileNotFoundError:
            logger.info("No cache file found at %s", CACHE_FILE)
        except Exception as e:
            logger.warning("Failed to load cache: %s", e)

    def _save_cache(self, current_data, known_models, loaded):
        """Save current scrape state to disk."""
        try:
            state = {
                "timestamp": time.time(),
                "known": list(known_models),
                "loaded": list(loaded),
                "models": {},
            }
            for mid, metrics in current_data.items():
                state["models"][mid] = dict(metrics)
            tmp = CACHE_FILE + ".tmp"
            with open(tmp, "w") as f:
                json.dump(state, f)
            os.replace(tmp, CACHE_FILE)
        except Exception as e:
            logger.error("Failed to save cache: %s", e)

    def _compute_deltas_for_data(self, current_data):
        """Compute deltas given current data and stored previous state.

        Only counter metrics get deltas. On counter reset (value went backward),
        delta is silently 0 (no entry added).
        """
        deltas = {}
        if not self._prev:
            return deltas

        for mid, prev_metrics in self._prev.items():
            mid_deltas = {}
            cur_metrics = current_data.get(mid, {})
            for mname, prev_val in prev_metrics.items():
                # Only compute deltas for counter metrics
                if mname not in COUNTER_METRICS:
                    continue
                cur_val = cur_metrics.get(mname)
                if cur_val is None:
                    continue
                diff = cur_val - prev_val
                # Counter reset: value went backward, delta is 0 (wrapped)
                if diff < 0:
                    continue
                if diff > 0:
                    mid_deltas[mname] = diff
            if mid_deltas:
                deltas[mid] = mid_deltas
        return deltas

    def update(self, models_data, metric_lines, known_models=None):
        """Called by the background thread after a successful scrape cycle."""
        with self._lock:
            # Discover model IDs and loaded status from /models endpoint
            known = set()
            loaded = {}
            if models_data and isinstance(models_data, dict):
                model_list = models_data.get("data", [])
                for m in model_list:
                    if not isinstance(m, dict):
                        continue
                    mid = m.get("id", "unknown")
                    known.add(mid)
                    status_data = m.get("status", {})
                    if isinstance(status_data, dict) and status_data.get("value") == "loaded":
                        loaded[mid] = 1.0
                    self._health = {"status": "ok", "model": mid}

            if not known:
                self._health = {"status": "error", "model": "unknown"}

            # Build new metrics dict for all known models
            new_data = {}
            for mid in known:
                new_data[mid] = {}

            # Apply parsed metric values from /metrics?model=<id>
            for metric_name, model_id, value in metric_lines:
                if model_id in known:
                    new_data[model_id][metric_name] = value

            # Ensure every known model has all ALL_METRICS entries
            for mid in known:
                for mname in ALL_METRICS:
                    if mname not in new_data[mid]:
                        new_data[mid][mname] = 0.0

            # Previously known models no longer in the list get zeroed out
            for mid in self._known_models - known:
                new_data[mid] = {m: 0.0 for m in ALL_METRICS}

            # Compute deltas before updating previous state
            deltas = self._compute_deltas_for_data(new_data)

            # Save previous state for next cycle
            self._save_cache(new_data, known, loaded)

            # Update state
            self._prev = {mid: dict(metrics) for mid, metrics in new_data.items()}
            self._known_models = known
            self._data = new_data
            self._loaded = loaded
            self._deltas = deltas

    def snapshot(self):
        """Return a frozen copy of the current cache state including deltas."""
        with self._lock:
            return {
                "data": {k: dict(v) for k, v in self._data.items()},
                "known": set(self._known_models),
                "health": dict(self._health) if self._health else {"status": "error", "model": "unknown"},
                "loaded": dict(self._loaded),
                "deltas": {k: dict(v) for k, v in self._deltas.items()},
            }


cache = Cache()


def _fetch_json(url, timeout=SCRAPE_TIMEOUT):
    try:
        req = urllib.request.Request(url, headers={"Accept": "application/json"})
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except Exception as e:
        logger.debug("Failed to fetch %s: %s", url, e)
        return None


def _fetch_text(url, timeout=SCRAPE_TIMEOUT):
    try:
        req = urllib.request.Request(url, headers={"Accept": "text/plain"})
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.read().decode("utf-8")
    except Exception as e:
        logger.debug("Failed to fetch %s: %s", url, e)
        return None


def _scrape_cycle():
    """One full scrape cycle: discover models, then scrape metrics per model."""
    targets = []
    env_targets = os.environ.get("LLAMA_TARGETS", "")
    if env_targets:
        try:
            targets = json.loads(env_targets)
        except json.JSONDecodeError:
            targets = []
    if not targets:
        targets = LLAMA_CPP_DEFAULTS

    all_metric_lines = []
    models_data = None
    known_models = {}  # { model_id: base_url, ... }

    for target in targets:
        base_url = target["url"].rstrip("/")

        # Fetch /models to discover models and their status
        models_data = _fetch_json(f"{base_url}/models")

        if models_data and isinstance(models_data, dict):
            model_list = models_data.get("data", [])
            for m in model_list:
                if not isinstance(m, dict) or "id" not in m:
                    continue
                model_id = m["id"]
                known_models[model_id] = base_url
                # Only scrape metrics from loaded models to avoid triggering loads
                status_data = m.get("status", {})
                if isinstance(status_data, dict) and status_data.get("value") != "loaded":
                    continue
                # Scrape /metrics for loaded models only
                metrics_url = f"{base_url}/metrics?model={model_id}"
                body = _fetch_text(metrics_url)
                if body:
                    for line in body.splitlines():
                        parsed = _parse_metric_line(line)
                        if parsed:
                            metric_name, metric_value = parsed
                            if metric_name in ALL_METRICS:
                                all_metric_lines.append((metric_name, model_id, metric_value))
                else:
                    logger.debug("No metrics body for model %s", model_id)

    # Update the shared cache
    cache.update(models_data, all_metric_lines, known_models)


def _parse_metric_line(line):
    """Parse a single Prometheus metric line. Returns (name, value) or None."""
    line = line.strip()
    if not line or line.startswith("#"):
        return None
    try:
        # Handle lines with labels: metric_name{label="val"} value
        if "{" in line:
            parts = line.split("{")
            name = parts[0].strip()
            rest = parts[1]
            # value is the last token after the closing }
            value = rest.rsplit("}", 1)[-1].strip().split()[-1] if "}" in rest else rest.strip()
        else:
            parts = line.split()
            name = parts[0]
            value = parts[1] if len(parts) >= 2 else "1"
        # Try to parse as float
        float(value)
        return (name, float(value))
    except (ValueError, IndexError):
        return None


def _background_scrape():
    """Background thread: periodically scrape and update cache."""
    logger.info("Background scraper started (interval=%ds)", POLL_INTERVAL)
    # Do one immediate scrape on startup
    _scrape_cycle()
    while True:
        try:
            time.sleep(POLL_INTERVAL)
            _scrape_cycle()
        except Exception as e:
            logger.error("Scrape cycle error: %s", e)


class MetricsHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path != "/metrics":
            self.send_response(404)
            self.end_headers()
            return

        snap = cache.snapshot()
        lines = []

        def _fmt(metric_name, model_id, value):
            return metric_name + '{' + 'server="llama-cpp-11434",model="' + model_id + '"} ' + str(value)

        # Health metric
        health = snap["health"]
        status = health.get("status", "error")
        health_model = health.get("model", "unknown")
        health_val = 1.0 if status == "ok" else 0.0
        lines.append(_fmt("llama_server_health", health_model, health_val))

        # Loaded metrics
        for mid in snap["loaded"]:
            lines.append(_fmt("llama_models_loaded", mid, snap["loaded"][mid]))

        # Per-model metrics from cache (absolute values)
        for mid in sorted(snap["data"]):
            metrics = snap["data"][mid]
            for mname in ALL_METRICS:
                value = metrics.get(mname, 0.0)
                lines.append(_fmt(mname, mid, value))

        # Per-model delta metrics (counters as change since last scrape)
        deltas = snap.get("deltas", {})
        for mid in sorted(deltas):
            delta_metrics = deltas[mid]
            for mname in sorted(delta_metrics):
                value = delta_metrics[mname]
                lines.append(_fmt(mname + "_delta", mid, value))

        body = "\n".join(lines) + "\n" if lines else "# no metrics available\n"

        self.send_response(200)
        self.send_header("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
        self.end_headers()
        self.wfile.write(body.encode("utf-8"))

    def log_message(self, format, *args):
        logger.debug("%s - - %s", self.address_string(), format % args)


def main():
    port = int(os.environ.get("LLAMA_EXPORTER_PORT", "9550"))
    host = os.environ.get("LLAMA_EXPORTER_BIND", "0.0.0.0")

    # Start background scraper thread
    scraper = threading.Thread(target=_background_scrape, daemon=True)
    scraper.start()

    # Start HTTP server
    server = HTTPServer((host, port), MetricsHandler)
    logger.info("Starting Llama Exporter on %s:%d (interval=%ds, timeout=%ds, cache=%s)",
                host, port, POLL_INTERVAL, SCRAPE_TIMEOUT, CACHE_FILE)

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        logger.info("Shutting down")
        server.shutdown()


if __name__ == "__main__":
    main()