compute/src/telemetry.rs

//! Privacy-safe, durable model-serving telemetry.
//!
//! The model runtime remains the authority for scheduling, token, latency, and
//! multimodal encoder metrics. Compute samples its Prometheus endpoint and stores
//! only numeric measurements plus bounded labels. Prompts, outputs, image data,
//! URLs, and arbitrary HTTP headers never enter this database.

use anyhow::{bail, Context, Result};
use rusqlite::{params, Connection};
use serde::Serialize;
use std::collections::{BTreeMap, HashMap};
use std::fs;
use std::io::{Read, Write};
use std::net::{TcpStream, ToSocketAddrs};
use std::path::{Path, PathBuf};
use std::thread;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use crate::config::{find_scene, Registry};
use crate::proc::State;

const DB_RELATIVE_PATH: &str = ".compute/usage/telemetry.sqlite3";
const MAX_METRICS_RESPONSE_BYTES: usize = 16 * 1024 * 1024;
const COUNTER_SNAPSHOT_INTERVAL_MS: i64 = 60_000;
const ALLOWED_CUSTOM_LABELS: &[&str] = &["client", "agent", "workload"];

const COUNTERS: &[&str] = &[
    "sglang:prompt_tokens_total",
    "sglang:generation_tokens_total",
    "sglang:cached_tokens_total",
    "sglang:num_requests_total",
    "sglang:num_aborted_requests_total",
    "sglang:encoder_requests_received_total",
    "sglang:encoder_cache_total_tokens_total",
    "sglang:encoder_cache_hit_tokens_total",
    "sglang:encoder_mm_items_per_request_sum",
    "sglang:encoder_mm_items_per_request_count",
    "sglang:time_to_first_token_seconds_sum",
    "sglang:time_to_first_token_seconds_count",
    "sglang:e2e_request_latency_seconds_sum",
    "sglang:e2e_request_latency_seconds_count",
];

#[derive(Debug, Clone, PartialEq)]
struct Metric {
    name: String,
    labels: BTreeMap<String, String>,
    value: f64,
}

#[derive(Debug, Clone)]
pub struct Target {
    pub model: String,
    pub scene: String,
    pub max_concurrency: u32,
    pub config_fingerprint: String,
}

#[derive(Debug, Serialize)]
pub struct CollectResult {
    pub timestamp_ms: i64,
    pub model: String,
    pub scene: String,
    pub max_concurrency: u32,
    pub metrics_recorded: usize,
    pub running: u32,
    pub queued: u32,
}

#[derive(Debug, Clone, Serialize, Default)]
pub struct UsageTotals {
    pub requests: f64,
    pub aborted_requests: f64,
    pub prompt_tokens: f64,
    pub generation_tokens: f64,
    pub cached_tokens: f64,
    pub vision_requests: f64,
    pub vision_request_share: f64,
    pub vision_request_source: String,
    pub vision_prompt_tokens: f64,
    pub vision_items: f64,
    pub vision_encoder_tokens: f64,
    pub vision_encoder_cached_tokens: f64,
}

#[derive(Debug, Clone, Serialize)]
pub struct ConcurrencyBucket {
    pub concurrency: String,
    pub samples: u64,
    pub share: f64,
}

#[derive(Debug, Clone, Serialize)]
pub struct ConcurrencySummary {
    pub samples: u64,
    pub queued_samples: u64,
    pub peak_running: u32,
    pub peak_queued: u32,
    /// A time-sampled distribution, not an admission-count distribution.
    pub distribution: Vec<ConcurrencyBucket>,
}

#[derive(Debug, Clone, Serialize)]
pub struct AgentUsage {
    pub client: String,
    pub agent: String,
    pub workload: String,
    pub requests: f64,
    pub prompt_tokens: f64,
    pub generation_tokens: f64,
}

#[derive(Debug, Clone, Serialize)]
pub struct UsageSummary {
    pub since_ms: i64,
    pub until_ms: i64,
    pub database: String,
    pub totals: UsageTotals,
    pub concurrency: ConcurrencySummary,
    pub agents: Vec<AgentUsage>,
}

pub struct Store {
    path: PathBuf,
    conn: Connection,
}

impl Store {
    pub fn open(root: &Path) -> Result<Self> {
        let path = root.join(DB_RELATIVE_PATH);
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)
                .with_context(|| format!("creating telemetry directory {}", parent.display()))?;
        }
        let conn = Connection::open(&path)
            .with_context(|| format!("opening telemetry database {}", path.display()))?;
        conn.busy_timeout(Duration::from_secs(5))?;
        conn.execute_batch(
            "PRAGMA journal_mode=WAL;
             PRAGMA synchronous=NORMAL;
             CREATE TABLE IF NOT EXISTS telemetry_schema (
               version INTEGER PRIMARY KEY,
               applied_at_ms INTEGER NOT NULL
             );
             CREATE TABLE IF NOT EXISTS config_snapshots (
               fingerprint TEXT PRIMARY KEY,
               first_seen_ms INTEGER NOT NULL,
               model TEXT NOT NULL,
               scene TEXT NOT NULL,
               max_concurrency INTEGER NOT NULL
             );
             CREATE TABLE IF NOT EXISTS engine_samples (
               timestamp_ms INTEGER NOT NULL,
               model TEXT NOT NULL,
               scene TEXT NOT NULL,
               config_fingerprint TEXT NOT NULL,
               max_concurrency INTEGER NOT NULL,
               running REAL NOT NULL,
               queued REAL NOT NULL,
               generation_tps REAL,
               cache_hit_rate REAL,
               PRIMARY KEY (timestamp_ms, model)
             );
             CREATE INDEX IF NOT EXISTS engine_samples_time
               ON engine_samples(timestamp_ms);
             CREATE TABLE IF NOT EXISTS counter_samples (
               timestamp_ms INTEGER NOT NULL,
               model TEXT NOT NULL,
               metric TEXT NOT NULL,
               labels_json TEXT NOT NULL,
               value REAL NOT NULL,
               PRIMARY KEY (timestamp_ms, model, metric, labels_json)
             );
             CREATE INDEX IF NOT EXISTS counter_samples_metric_time
               ON counter_samples(metric, timestamp_ms);
             INSERT OR IGNORE INTO telemetry_schema(version, applied_at_ms)
               VALUES (1, CAST(strftime('%s','now') AS INTEGER) * 1000);",
        )?;
        Ok(Self { path, conn })
    }

    fn record(&mut self, target: &Target, metrics: &[Metric]) -> Result<CollectResult> {
        self.record_with_counter_interval(target, metrics, COUNTER_SNAPSHOT_INTERVAL_MS)
    }

    fn record_with_counter_interval(
        &mut self,
        target: &Target,
        metrics: &[Metric],
        counter_interval_ms: i64,
    ) -> Result<CollectResult> {
        let timestamp_ms = unix_time_ms()?;
        let running = default_gauge(metrics, "sglang:num_running_reqs").unwrap_or(0.0);
        let queued = default_gauge(metrics, "sglang:num_queue_reqs").unwrap_or(0.0);
        let generation_tps = default_gauge(metrics, "sglang:gen_throughput");
        let cache_hit_rate = default_gauge(metrics, "sglang:cache_hit_rate");

        let last_counter_ms: Option<i64> = self.conn.query_row(
            "SELECT MAX(timestamp_ms) FROM counter_samples WHERE model = ?1",
            params![target.model],
            |row| row.get(0),
        )?;
        let counter_due =
            last_counter_ms.is_none_or(|previous| timestamp_ms - previous >= counter_interval_ms);

        let transaction = self.conn.transaction()?;
        transaction.execute(
            "INSERT OR IGNORE INTO config_snapshots
             (fingerprint, first_seen_ms, model, scene, max_concurrency)
             VALUES (?1, ?2, ?3, ?4, ?5)",
            params![
                target.config_fingerprint,
                timestamp_ms,
                target.model,
                target.scene,
                target.max_concurrency
            ],
        )?;
        transaction.execute(
            "INSERT INTO engine_samples
             (timestamp_ms, model, scene, config_fingerprint, max_concurrency,
              running, queued, generation_tps, cache_hit_rate)
             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
            params![
                timestamp_ms,
                target.model,
                target.scene,
                target.config_fingerprint,
                target.max_concurrency,
                running,
                queued,
                generation_tps,
                cache_hit_rate
            ],
        )?;

        let mut recorded = 0usize;
        if counter_due {
            for metric in metrics
                .iter()
                .filter(|metric| COUNTERS.contains(&metric.name.as_str()))
            {
                let labels = retained_labels(&metric.labels);
                transaction.execute(
                    "INSERT INTO counter_samples
                     (timestamp_ms, model, metric, labels_json, value)
                     VALUES (?1, ?2, ?3, ?4, ?5)",
                    params![
                        timestamp_ms,
                        target.model,
                        metric.name,
                        serde_json::to_string(&labels)?,
                        metric.value
                    ],
                )?;
                recorded += 1;
            }
        }
        transaction.commit()?;

        Ok(CollectResult {
            timestamp_ms,
            model: target.model.clone(),
            scene: target.scene.clone(),
            max_concurrency: target.max_concurrency,
            metrics_recorded: recorded,
            running: running.max(0.0).round() as u32,
            queued: queued.max(0.0).round() as u32,
        })
    }

    pub fn summary(&self, since_ms: i64, until_ms: i64) -> Result<UsageSummary> {
        let image = |labels: &BTreeMap<String, String>| {
            labels.get("modality").is_some_and(|value| value == "image")
        };
        let vision_workload = |labels: &BTreeMap<String, String>| {
            labels
                .get("workload")
                .is_some_and(|value| value == "vision")
        };
        let has_encoder_request_metrics =
            self.has_metric_samples("sglang:encoder_requests_received_total", since_ms, until_ms)?;
        let vision_requests = if has_encoder_request_metrics {
            self.metric_delta(
                "sglang:encoder_requests_received_total",
                since_ms,
                until_ms,
                image,
            )?
        } else {
            self.metric_delta(
                "sglang:num_requests_total",
                since_ms,
                until_ms,
                vision_workload,
            )?
        };
        let mut totals = UsageTotals {
            requests: self
                .metric_delta("sglang:num_requests_total", since_ms, until_ms, |_| true)?,
            aborted_requests: self.metric_delta(
                "sglang:num_aborted_requests_total",
                since_ms,
                until_ms,
                |_| true,
            )?,
            prompt_tokens: self.metric_delta(
                "sglang:prompt_tokens_total",
                since_ms,
                until_ms,
                |_| true,
            )?,
            generation_tokens: self.metric_delta(
                "sglang:generation_tokens_total",
                since_ms,
                until_ms,
                |_| true,
            )?,
            cached_tokens: self.metric_delta(
                "sglang:cached_tokens_total",
                since_ms,
                until_ms,
                |_| true,
            )?,
            vision_requests,
            vision_request_share: 0.0,
            vision_request_source: if has_encoder_request_metrics {
                "encoder_metrics".to_string()
            } else {
                "workload_label".to_string()
            },
            vision_prompt_tokens: self.metric_delta(
                "sglang:prompt_tokens_total",
                since_ms,
                until_ms,
                vision_workload,
            )?,
            vision_items: self.metric_delta(
                "sglang:encoder_mm_items_per_request_sum",
                since_ms,
                until_ms,
                image,
            )?,
            vision_encoder_tokens: self.metric_delta(
                "sglang:encoder_cache_total_tokens_total",
                since_ms,
                until_ms,
                image,
            )?,
            vision_encoder_cached_tokens: self.metric_delta(
                "sglang:encoder_cache_hit_tokens_total",
                since_ms,
                until_ms,
                image,
            )?,
        };
        totals.vision_request_share = if totals.requests > 0.0 {
            (totals.vision_requests / totals.requests).clamp(0.0, 1.0)
        } else {
            0.0
        };

        Ok(UsageSummary {
            since_ms,
            until_ms,
            database: self.path.display().to_string(),
            totals,
            concurrency: self.concurrency(since_ms, until_ms)?,
            agents: self.agent_usage(since_ms, until_ms)?,
        })
    }

    fn metric_delta<F>(&self, metric: &str, since_ms: i64, until_ms: i64, keep: F) -> Result<f64>
    where
        F: Fn(&BTreeMap<String, String>) -> bool,
    {
        let total: f64 = self
            .series_deltas(metric, since_ms, until_ms)?
            .into_iter()
            .filter(|(labels, _)| keep(labels))
            .map(|(_, delta)| delta)
            .sum();
        Ok(if total == 0.0 { 0.0 } else { total })
    }

    fn has_metric_samples(&self, metric: &str, since_ms: i64, until_ms: i64) -> Result<bool> {
        self.conn
            .query_row(
                "SELECT EXISTS(
                   SELECT 1 FROM counter_samples
                   WHERE metric = ?1 AND timestamp_ms >= ?2 AND timestamp_ms <= ?3
                 )",
                params![metric, since_ms, until_ms],
                |row| row.get(0),
            )
            .context("checking telemetry metric availability")
    }

    fn series_deltas(
        &self,
        metric: &str,
        since_ms: i64,
        until_ms: i64,
    ) -> Result<Vec<(BTreeMap<String, String>, f64)>> {
        let mut baseline_statement = self.conn.prepare(
            "SELECT model, MIN(timestamp_ms)
             FROM counter_samples
             WHERE metric = ?1 AND timestamp_ms <= ?2
             GROUP BY model",
        )?;
        let baseline_rows = baseline_statement.query_map(params![metric, until_ms], |row| {
            Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?))
        })?;
        let mut model_baselines = HashMap::new();
        for row in baseline_rows {
            let (model, timestamp_ms) = row?;
            model_baselines.insert(model, timestamp_ms);
        }

        let mut statement = self.conn.prepare(
            "SELECT model, timestamp_ms, labels_json, value
             FROM counter_samples
             WHERE metric = ?1 AND timestamp_ms <= ?2
             ORDER BY model, labels_json, timestamp_ms",
        )?;
        let rows = statement.query_map(params![metric, until_ms], |row| {
            Ok((
                row.get::<_, String>(0)?,
                row.get::<_, i64>(1)?,
                row.get::<_, String>(2)?,
                row.get::<_, f64>(3)?,
            ))
        })?;

        let mut state: HashMap<(String, String), (Option<f64>, f64)> = HashMap::new();
        for row in rows {
            let (model, timestamp_ms, labels_json, value) = row?;
            let model_baseline = model_baselines.get(&model).copied().unwrap_or(timestamp_ms);
            let entry = state.entry((model, labels_json)).or_insert((None, 0.0));
            if timestamp_ms >= since_ms {
                if let Some(previous) = entry.0 {
                    entry.1 += if value >= previous {
                        value - previous
                    } else {
                        value
                    };
                } else if timestamp_ms > model_baseline {
                    // Prometheus creates a new labeled series on first use. The
                    // model-wide baseline predates it, so its initial value is
                    // usage in this window rather than unknown prehistory.
                    entry.1 += value;
                }
            }
            entry.0 = Some(value);
        }

        state
            .into_iter()
            .map(|((_model, labels), (_previous, delta))| {
                let labels = serde_json::from_str(&labels)
                    .context("parsing labels stored in telemetry database")?;
                Ok((labels, delta))
            })
            .collect()
    }

    fn concurrency(&self, since_ms: i64, until_ms: i64) -> Result<ConcurrencySummary> {
        let mut statement = self.conn.prepare(
            "SELECT running, queued FROM engine_samples
             WHERE timestamp_ms >= ?1 AND timestamp_ms <= ?2
             ORDER BY timestamp_ms",
        )?;
        let rows = statement.query_map(params![since_ms, until_ms], |row| {
            Ok((row.get::<_, f64>(0)?, row.get::<_, f64>(1)?))
        })?;
        let mut counts: BTreeMap<String, u64> = BTreeMap::new();
        let mut samples = 0u64;
        let mut queued_samples = 0u64;
        let mut peak_running = 0u32;
        let mut peak_queued = 0u32;
        for row in rows {
            let (running, queued) = row?;
            let running = running.max(0.0).round() as u32;
            let queued = queued.max(0.0).round() as u32;
            let bucket = if running >= 4 {
                "C4+".to_string()
            } else {
                format!("C{running}")
            };
            *counts.entry(bucket).or_default() += 1;
            samples += 1;
            queued_samples += u64::from(queued > 0);
            peak_running = peak_running.max(running);
            peak_queued = peak_queued.max(queued);
        }
        let distribution = counts
            .into_iter()
            .map(|(concurrency, count)| ConcurrencyBucket {
                concurrency,
                samples: count,
                share: if samples == 0 {
                    0.0
                } else {
                    count as f64 / samples as f64
                },
            })
            .collect();
        Ok(ConcurrencySummary {
            samples,
            queued_samples,
            peak_running,
            peak_queued,
            distribution,
        })
    }

    fn agent_usage(&self, since_ms: i64, until_ms: i64) -> Result<Vec<AgentUsage>> {
        #[derive(Default)]
        struct Row {
            requests: f64,
            prompt_tokens: f64,
            generation_tokens: f64,
        }
        let mut grouped: BTreeMap<(String, String, String), Row> = BTreeMap::new();
        for (metric, field) in [
            ("sglang:num_requests_total", 0),
            ("sglang:prompt_tokens_total", 1),
            ("sglang:generation_tokens_total", 2),
        ] {
            for (labels, delta) in self.series_deltas(metric, since_ms, until_ms)? {
                let key = (
                    label_or_unknown(&labels, "client"),
                    label_or_unknown(&labels, "agent"),
                    label_or_unknown(&labels, "workload"),
                );
                let row = grouped.entry(key).or_default();
                match field {
                    0 => row.requests += delta,
                    1 => row.prompt_tokens += delta,
                    _ => row.generation_tokens += delta,
                }
            }
        }
        let mut rows: Vec<_> = grouped
            .into_iter()
            .filter(|(_, row)| {
                row.requests != 0.0 || row.prompt_tokens != 0.0 || row.generation_tokens != 0.0
            })
            .map(|((client, agent, workload), row)| AgentUsage {
                client,
                agent,
                workload,
                requests: row.requests,
                prompt_tokens: row.prompt_tokens,
                generation_tokens: row.generation_tokens,
            })
            .collect();
        rows.sort_by(|a, b| b.requests.total_cmp(&a.requests));
        Ok(rows)
    }
}

pub fn collect_once(root: &Path, metrics_url: &str, model_hint: &str) -> Result<CollectResult> {
    let target = resolve_target(root, metrics_url, model_hint)?;
    let body = fetch_http(metrics_url)?;
    let metrics = parse_prometheus(&body);
    if !metrics
        .iter()
        .any(|metric| metric.name == "sglang:num_running_reqs")
    {
        bail!("{metrics_url} returned no SGLang scheduler metrics; start SGLang with --enable-metrics");
    }
    Store::open(root)?.record(&target, &metrics)
}

pub fn run_collector(root: PathBuf, metrics_url: String, model_hint: String, interval: Duration) {
    let mut failures = 0u64;
    loop {
        match collect_once(&root, &metrics_url, &model_hint) {
            Ok(result) => {
                if failures > 0 {
                    eprintln!(
                        "telemetry: recovered after {failures} failed scrape(s); model={} scene={}",
                        result.model, result.scene
                    );
                }
                failures = 0;
            }
            Err(error) => {
                failures += 1;
                if failures == 1 || failures.is_multiple_of(60) {
                    eprintln!("telemetry: scrape failed ({failures} consecutive): {error:#}");
                }
            }
        }
        thread::sleep(interval.max(Duration::from_secs(1)));
    }
}

pub fn report(root: &Path, since: &str) -> Result<UsageSummary> {
    let duration = parse_duration(since)?;
    let until_ms = unix_time_ms()?;
    let since_ms = until_ms.saturating_sub(duration.as_millis().min(i64::MAX as u128) as i64);
    Store::open(root)?.summary(since_ms, until_ms)
}

pub fn render_text(summary: &UsageSummary) -> String {
    let totals = &summary.totals;
    let mut lines = vec![
        "Lumbridge Compute · usage".to_string(),
        format!("  requests     {:>12.0}", totals.requests),
        format!("  prompt tok   {:>12.0}", totals.prompt_tokens),
        format!("  output tok   {:>12.0}", totals.generation_tokens),
        format!("  cached tok   {:>12.0}", totals.cached_tokens),
        format!("  aborted      {:>12.0}", totals.aborted_requests),
        format!("  vision calls {:>12.0}", totals.vision_requests),
        format!(
            "  vision share {:>11.1}%",
            totals.vision_request_share * 100.0
        ),
        format!("  vision tok   {:>12.0}", totals.vision_prompt_tokens),
        format!("  image items  {:>12.0}", totals.vision_items),
        String::new(),
        "  time-sampled concurrency:".to_string(),
    ];
    if summary.concurrency.samples == 0 {
        lines.push("    (no samples in this window)".to_string());
    } else {
        for bucket in &summary.concurrency.distribution {
            lines.push(format!(
                "    {:>3}  {:>8} samples  {:>6.1}%",
                bucket.concurrency,
                bucket.samples,
                bucket.share * 100.0
            ));
        }
        lines.push(format!(
            "    peak running {} · peak queued {} · queue observed in {} sample(s)",
            summary.concurrency.peak_running,
            summary.concurrency.peak_queued,
            summary.concurrency.queued_samples
        ));
    }
    if !summary.agents.is_empty() {
        lines.push(String::new());
        lines.push("  agents:".to_string());
        for row in &summary.agents {
            lines.push(format!(
                "    {}/{}/{}  {:.0} req · {:.0} in · {:.0} out tok",
                row.client,
                row.agent,
                row.workload,
                row.requests,
                row.prompt_tokens,
                row.generation_tokens
            ));
        }
    }
    lines.join("\n")
}

fn resolve_target(root: &Path, metrics_url: &str, model_hint: &str) -> Result<Target> {
    let registry = Registry::load(root)?;
    let state = State::load_checked(root)?;
    let scene_name = state
        .active_scene
        .or(state.desired_scene)
        .unwrap_or_else(|| "unknown".to_string());
    let port = parse_http_url(metrics_url)?.1;
    let model = if model_hint != "auto" {
        model_hint.to_string()
    } else if scene_name != "unknown" {
        find_scene(root, &scene_name)?
            .models
            .into_iter()
            .find(|id| registry.models.get(id).and_then(|m| m.serve.port) == Some(port))
            .with_context(|| {
                format!("active Scene '{scene_name}' has no model on metrics port {port}")
            })?
    } else {
        registry
            .models
            .iter()
            .find(|(_, model)| model.serve.port == Some(port))
            .map(|(id, _)| id.clone())
            .with_context(|| format!("registry has no model on metrics port {port}"))?
    };
    let model_config = registry
        .models
        .get(&model)
        .with_context(|| format!("no registry model named '{model}'"))?;
    let max_concurrency = configured_concurrency(model_config);
    let fingerprint_input = format!(
        "{model}\n{scene_name}\n{max_concurrency}\n{:?}\n{:?}",
        model_config.serve.command, model_config.serve.args
    );
    Ok(Target {
        model,
        scene: scene_name,
        max_concurrency,
        config_fingerprint: format!("fnv1a64:{:016x}", fnv1a64(fingerprint_input.as_bytes())),
    })
}

fn configured_concurrency(model: &crate::config::Model) -> u32 {
    for flag in ["--max-running-requests", "--max-num-seqs"] {
        if let Some(index) = model.serve.command.iter().position(|arg| arg == flag) {
            if let Some(value) = model
                .serve
                .command
                .get(index + 1)
                .and_then(|value| value.parse().ok())
            {
                return value;
            }
        }
    }
    for key in ["max-running-requests", "max-num-seqs"] {
        if let Some(value) = model
            .serve
            .args
            .get(key)
            .and_then(serde_yaml::Value::as_u64)
        {
            return value.min(u32::MAX as u64) as u32;
        }
    }
    0
}

fn retained_labels(labels: &BTreeMap<String, String>) -> BTreeMap<String, String> {
    labels
        .iter()
        .filter(|(key, _)| {
            matches!(
                key.as_str(),
                "model_name" | "stream" | "reason" | "modality" | "status" | "priority"
            ) || ALLOWED_CUSTOM_LABELS.contains(&key.as_str())
        })
        .map(|(key, value)| {
            let value = if ALLOWED_CUSTOM_LABELS.contains(&key.as_str()) {
                sanitize_custom_label(value)
            } else {
                value.chars().take(128).collect()
            };
            (key.clone(), value)
        })
        .collect()
}

fn sanitize_custom_label(value: &str) -> String {
    if !value.is_empty()
        && value.len() <= 64
        && value
            .bytes()
            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-'))
    {
        value.to_string()
    } else {
        "invalid".to_string()
    }
}

fn label_or_unknown(labels: &BTreeMap<String, String>, key: &str) -> String {
    labels
        .get(key)
        // SGLang emits empty custom labels on unlabeled traffic. They are
        // deliberately collapsed to the same bounded sentinel as malformed
        // values at ingestion, and both render as unattributed usage.
        .filter(|value| !value.is_empty() && value.as_str() != "invalid")
        .cloned()
        .unwrap_or_else(|| "unknown".to_string())
}

fn default_gauge(metrics: &[Metric], name: &str) -> Option<f64> {
    metrics
        .iter()
        .filter(|metric| metric.name == name)
        .find(|metric| metric.labels.get("priority").is_none_or(String::is_empty))
        .map(|metric| metric.value)
}

fn parse_prometheus(text: &str) -> Vec<Metric> {
    text.lines().filter_map(parse_metric_line).collect()
}

fn parse_metric_line(line: &str) -> Option<Metric> {
    let line = line.trim();
    if line.is_empty() || line.starts_with('#') {
        return None;
    }
    let split = line.rfind(char::is_whitespace)?;
    let head = line[..split].trim();
    let value = line[split..].trim().parse::<f64>().ok()?;
    if !value.is_finite() {
        return None;
    }
    let (name, labels) = if let Some(open) = head.find('{') {
        if !head.ends_with('}') {
            return None;
        }
        (
            &head[..open],
            parse_labels(&head[open + 1..head.len() - 1])?,
        )
    } else {
        (head, BTreeMap::new())
    };
    Some(Metric {
        name: name.to_string(),
        labels,
        value,
    })
}

fn parse_labels(input: &str) -> Option<BTreeMap<String, String>> {
    let bytes = input.as_bytes();
    let mut labels = BTreeMap::new();
    let mut index = 0usize;
    while index < bytes.len() {
        while index < bytes.len() && (bytes[index] == b',' || bytes[index].is_ascii_whitespace()) {
            index += 1;
        }
        if index == bytes.len() {
            break;
        }
        let key_start = index;
        while index < bytes.len() && bytes[index] != b'=' {
            index += 1;
        }
        if index == bytes.len() {
            return None;
        }
        let key = input[key_start..index].trim().to_string();
        index += 1;
        if bytes.get(index) != Some(&b'"') {
            return None;
        }
        index += 1;
        let mut value = String::new();
        while index < bytes.len() {
            match bytes[index] {
                b'"' => {
                    index += 1;
                    break;
                }
                b'\\' => {
                    index += 1;
                    let escaped = *bytes.get(index)?;
                    value.push(match escaped {
                        b'n' => '\n',
                        b'\\' => '\\',
                        b'"' => '"',
                        other => other as char,
                    });
                    index += 1;
                }
                byte => {
                    value.push(byte as char);
                    index += 1;
                }
            }
        }
        labels.insert(key, value);
    }
    Some(labels)
}

fn fetch_http(url: &str) -> Result<String> {
    let (host, port, path) = parse_http_url(url)?;
    let address = (host.as_str(), port)
        .to_socket_addrs()?
        .next()
        .with_context(|| format!("resolving {host}:{port}"))?;
    let mut stream = TcpStream::connect_timeout(&address, Duration::from_secs(3))
        .with_context(|| format!("connecting to {url}"))?;
    stream.set_read_timeout(Some(Duration::from_secs(5)))?;
    stream.set_write_timeout(Some(Duration::from_secs(3)))?;
    write!(
        stream,
        "GET {path} HTTP/1.1\r\nHost: {host}:{port}\r\nAccept: text/plain\r\nAccept-Encoding: identity\r\nConnection: close\r\n\r\n"
    )?;
    let mut response = Vec::new();
    stream
        .take((MAX_METRICS_RESPONSE_BYTES + 1) as u64)
        .read_to_end(&mut response)?;
    if response.len() > MAX_METRICS_RESPONSE_BYTES {
        bail!("metrics response exceeded {MAX_METRICS_RESPONSE_BYTES} bytes");
    }
    let head_end = response
        .windows(4)
        .position(|window| window == b"\r\n\r\n")
        .context("metrics endpoint returned no HTTP header terminator")?;
    let head = String::from_utf8_lossy(&response[..head_end]);
    let status = head.lines().next().unwrap_or_default();
    if !status
        .split_whitespace()
        .nth(1)
        .is_some_and(|code| code == "200")
    {
        bail!("metrics endpoint returned {status}");
    }
    let body = &response[head_end + 4..];
    let body = if head
        .lines()
        .any(|line| line.eq_ignore_ascii_case("transfer-encoding: chunked"))
    {
        decode_chunked(body)?
    } else {
        body.to_vec()
    };
    String::from_utf8(body).context("metrics response was not UTF-8")
}

fn decode_chunked(mut body: &[u8]) -> Result<Vec<u8>> {
    let mut decoded = Vec::new();
    loop {
        let line_end = body
            .windows(2)
            .position(|window| window == b"\r\n")
            .context("invalid chunked metrics response")?;
        let size_text = std::str::from_utf8(&body[..line_end])?
            .split(';')
            .next()
            .unwrap_or_default();
        let size = usize::from_str_radix(size_text.trim(), 16)?;
        body = &body[line_end + 2..];
        if size == 0 {
            break;
        }
        if body.len() < size + 2 {
            bail!("truncated chunked metrics response");
        }
        decoded.extend_from_slice(&body[..size]);
        body = &body[size + 2..];
    }
    Ok(decoded)
}

fn parse_http_url(url: &str) -> Result<(String, u16, String)> {
    let rest = url
        .strip_prefix("http://")
        .with_context(|| format!("telemetry metrics URL must use local http://, got '{url}'"))?;
    let (authority, path) = rest.split_once('/').unwrap_or((rest, ""));
    let (host, port) = authority
        .rsplit_once(':')
        .with_context(|| format!("metrics URL must include a port: '{url}'"))?;
    let port = port.parse::<u16>().context("parsing metrics URL port")?;
    if host.is_empty() {
        bail!("metrics URL has an empty host");
    }
    Ok((host.to_string(), port, format!("/{path}")))
}

fn parse_duration(value: &str) -> Result<Duration> {
    let split = value
        .find(|ch: char| !ch.is_ascii_digit())
        .unwrap_or(value.len());
    let amount = value[..split]
        .parse::<u64>()
        .with_context(|| format!("invalid duration '{value}'"))?;
    let unit = &value[split..];
    let seconds = match unit {
        "s" => amount,
        "m" => amount.saturating_mul(60),
        "h" => amount.saturating_mul(60 * 60),
        "d" => amount.saturating_mul(24 * 60 * 60),
        "w" => amount.saturating_mul(7 * 24 * 60 * 60),
        _ => bail!("duration must end in s, m, h, d, or w (for example 24h or 7d)"),
    };
    Ok(Duration::from_secs(seconds))
}

fn unix_time_ms() -> Result<i64> {
    let milliseconds = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .context("system clock is before Unix epoch")?
        .as_millis();
    Ok(milliseconds.min(i64::MAX as u128) as i64)
}

fn fnv1a64(input: &[u8]) -> u64 {
    let mut hash = 0xcbf29ce484222325u64;
    for byte in input {
        hash ^= u64::from(*byte);
        hash = hash.wrapping_mul(0x100000001b3);
    }
    hash
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::{SystemTime, UNIX_EPOCH};

    fn temp_root() -> PathBuf {
        let suffix = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let path = std::env::temp_dir().join(format!("lumbridge-telemetry-{suffix}"));
        fs::create_dir_all(&path).unwrap();
        path
    }

    #[test]
    fn prometheus_parser_keeps_escaped_bounded_labels() {
        let metrics = parse_prometheus(
            "# HELP ignored\nsglang:num_requests_total{model_name=\"brain\",agent=\"prime\\\"one\",stream=\"true\"} 12\nsglang:num_running_reqs 3\n",
        );
        assert_eq!(metrics.len(), 2);
        assert_eq!(metrics[0].name, "sglang:num_requests_total");
        assert_eq!(metrics[0].labels["agent"], "prime\"one");
        assert_eq!(metrics[1].value, 3.0);
    }

    #[test]
    fn counter_deltas_survive_runtime_resets() {
        let root = temp_root();
        let mut store = Store::open(&root).unwrap();
        let target = Target {
            model: "brain".to_string(),
            scene: "solo".to_string(),
            max_concurrency: 4,
            config_fingerprint: "test".to_string(),
        };
        let labels = BTreeMap::from([
            ("agent".to_string(), "a".to_string()),
            ("client".to_string(), "c".to_string()),
        ]);
        let mut record = |value| {
            store
                .record_with_counter_interval(
                    &target,
                    &[
                        Metric {
                            name: "sglang:num_running_reqs".to_string(),
                            labels: BTreeMap::new(),
                            value: 1.0,
                        },
                        Metric {
                            name: "sglang:num_requests_total".to_string(),
                            labels: labels.clone(),
                            value,
                        },
                    ],
                    0,
                )
                .unwrap();
            thread::sleep(Duration::from_millis(2));
        };
        record(10.0);
        let since = unix_time_ms().unwrap() - 1;
        record(14.0);
        record(2.0); // runtime restarted
        let summary = store.summary(since, unix_time_ms().unwrap()).unwrap();
        assert_eq!(summary.totals.requests, 6.0);
        assert_eq!(summary.agents[0].requests, 6.0);
        fs::remove_dir_all(root).unwrap();
    }

    #[test]
    fn new_labeled_series_count_from_first_use_and_identify_vision() {
        let root = temp_root();
        let mut store = Store::open(&root).unwrap();
        let target = Target {
            model: "brain".to_string(),
            scene: "solo".to_string(),
            max_concurrency: 4,
            config_fingerprint: "test".to_string(),
        };
        let gauge = Metric {
            name: "sglang:num_running_reqs".to_string(),
            labels: BTreeMap::new(),
            value: 0.0,
        };
        let baseline_labels = BTreeMap::from([
            ("agent".to_string(), "".to_string()),
            ("client".to_string(), "".to_string()),
            ("workload".to_string(), "".to_string()),
        ]);
        store
            .record_with_counter_interval(
                &target,
                &[
                    gauge.clone(),
                    Metric {
                        name: "sglang:num_requests_total".to_string(),
                        labels: baseline_labels.clone(),
                        value: 10.0,
                    },
                    Metric {
                        name: "sglang:prompt_tokens_total".to_string(),
                        labels: baseline_labels,
                        value: 100.0,
                    },
                ],
                0,
            )
            .unwrap();
        thread::sleep(Duration::from_millis(2));
        let since = unix_time_ms().unwrap() - 1;
        let vision_labels = BTreeMap::from([
            ("agent".to_string(), "eye-1".to_string()),
            ("client".to_string(), "cloud-agent".to_string()),
            ("workload".to_string(), "vision".to_string()),
        ]);
        store
            .record_with_counter_interval(
                &target,
                &[
                    gauge,
                    Metric {
                        name: "sglang:num_requests_total".to_string(),
                        labels: vision_labels.clone(),
                        value: 1.0,
                    },
                    Metric {
                        name: "sglang:prompt_tokens_total".to_string(),
                        labels: vision_labels,
                        value: 334.0,
                    },
                ],
                0,
            )
            .unwrap();

        let summary = store.summary(since, unix_time_ms().unwrap()).unwrap();
        assert_eq!(summary.totals.requests, 1.0);
        assert_eq!(summary.totals.vision_requests, 1.0);
        assert_eq!(summary.totals.vision_prompt_tokens, 334.0);
        assert_eq!(summary.totals.vision_request_source, "workload_label");
        assert_eq!(summary.agents[0].agent, "eye-1");
        fs::remove_dir_all(root).unwrap();
    }

    #[test]
    fn duration_parser_is_explicit() {
        assert_eq!(parse_duration("24h").unwrap(), Duration::from_secs(86_400));
        assert_eq!(parse_duration("7d").unwrap(), Duration::from_secs(604_800));
        assert!(parse_duration("24").is_err());
    }

    #[test]
    fn caller_labels_are_bounded_and_cannot_carry_structured_data() {
        assert_eq!(sanitize_custom_label(""), "invalid");
        assert_eq!(sanitize_custom_label("worker-1"), "worker-1");
        assert_eq!(sanitize_custom_label("person@example.com"), "invalid");
        assert_eq!(sanitize_custom_label(&"x".repeat(65)), "invalid");
        assert_eq!(
            label_or_unknown(
                &BTreeMap::from([("agent".to_string(), "invalid".to_string())]),
                "agent"
            ),
            "unknown"
        );
    }
}