cortiq-gateway 0.2.43

//! Request statistics tracking for the dashboard and `GET /metrics`.
//!
//! Everything is held in memory (aggregates + per-minute buckets for the retention
//! window + a ring buffer of recent requests) and optionally appended to a JSONL
//! file that is replayed on startup — so statistics survive a restart with no
//! database dependency.

use crate::config::StatsCfg;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, VecDeque};
use std::io::Write;
use std::sync::Mutex;
use std::time::{SystemTime, UNIX_EPOCH};

fn now_secs() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

/// A single record for a completed request (one JSONL line and one "recent" ring entry).
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RequestRecord {
    pub ts: u64,
    #[serde(default)]
    pub account: String,
    #[serde(default)]
    pub protocol: String,
    #[serde(default)]
    pub directive: String, // auto | pinned
    #[serde(default)]
    pub task_label: String,
    #[serde(default)]
    pub tier: String,
    #[serde(default)]
    pub score: f32,
    #[serde(default)]
    pub model_id: String,
    #[serde(default)]
    pub route_source: String,
    #[serde(default)]
    pub prompt_tokens: u32,
    #[serde(default)]
    pub completion_tokens: u32,
    #[serde(default)]
    pub cost_usd: f64,
    #[serde(default)]
    pub latency_ms: u64,
    #[serde(default)]
    pub outcome: String, // ok | error
    #[serde(default)]
    pub failover: bool,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
}

#[derive(Clone, Debug, Default, Serialize, serde::Deserialize)]
pub struct Totals {
    pub requests: u64,
    pub ok: u64,
    pub errors: u64,
    pub failovers: u64,
    pub prompt_tokens: u64,
    pub completion_tokens: u64,
    pub cost_usd: f64,
    pub latency_ms_sum: u64,
}

impl Totals {
    fn apply(&mut self, r: &RequestRecord) {
        self.requests += 1;
        if r.outcome == "ok" {
            self.ok += 1;
        } else {
            self.errors += 1;
        }
        if r.failover {
            self.failovers += 1;
        }
        self.prompt_tokens += r.prompt_tokens as u64;
        self.completion_tokens += r.completion_tokens as u64;
        self.cost_usd += r.cost_usd;
        self.latency_ms_sum += r.latency_ms;
    }
    /// Accumulate another Totals into this one (for windowed aggregation).
    fn merge(&mut self, o: &Totals) {
        self.requests += o.requests;
        self.ok += o.ok;
        self.errors += o.errors;
        self.failovers += o.failovers;
        self.prompt_tokens += o.prompt_tokens;
        self.completion_tokens += o.completion_tokens;
        self.cost_usd += o.cost_usd;
        self.latency_ms_sum += o.latency_ms_sum;
    }
}

/// Per-minute bucket. Beyond the chart fields it carries the per-group breakdown
/// and failovers so the admin snapshot can be computed for ANY time window by
/// summing the in-window buckets (the internal fields are not serialized into
/// the `series` payload).
#[derive(Clone, Debug, Default, Serialize)]
pub struct Bucket {
    pub minute: u64, // unix seconds, rounded down to the minute
    pub requests: u64,
    pub errors: u64,
    pub prompt_tokens: u64,
    pub completion_tokens: u64,
    pub cost_usd: f64,
    pub latency_ms_sum: u64,
    #[serde(skip)]
    pub failovers: u64,
    #[serde(skip)]
    pub by_model: HashMap<String, Totals>,
    #[serde(skip)]
    pub by_account: HashMap<String, Totals>,
    #[serde(skip)]
    pub by_tier: HashMap<String, Totals>,
    #[serde(skip)]
    pub by_label: HashMap<String, Totals>,
}

#[derive(Default)]
struct Inner {
    /// All-time totals + per-model (used by the Prometheus counters, which must
    /// be monotonic — the windowed admin snapshot is derived from the buckets).
    total: Totals,
    by_model: HashMap<String, Totals>,
    /// Per-minute buckets — the single source for the windowed admin snapshot.
    buckets: VecDeque<Bucket>,
    recent: VecDeque<RequestRecord>,
    /// Per-account totals over EVERY record still in the JSONL file (no
    /// retention gate) — per-key accounting for the keys page.
    by_account_alltime: HashMap<String, Totals>,
    /// Per-account totals folded out of the file by compaction. alltime(key)
    /// = folded + by_account_alltime, with no double counting: a record is
    /// either in the file or folded, never both.
    folded: HashMap<String, Totals>,
    /// Month-to-date spend per account ("YYYY-MM", usd) — budgets check this.
    month_spend: HashMap<String, (String, f64)>,
    /// Records with ts ≤ this are already reflected in the persisted
    /// month_spend (watermark against double counting on replay).
    month_upto_ts: u64,
    /// Appends since the last compaction size check.
    writes_since_check: u32,
}

pub struct Stats {
    enabled: bool,
    file: Option<String>,
    ring_size: usize,
    retention_secs: u64,
    max_file_bytes: u64,
    inner: Mutex<Inner>,
}

impl Stats {
    pub fn new(cfg: &StatsCfg) -> std::sync::Arc<Self> {
        let file = if cfg.file.trim().is_empty() {
            None
        } else {
            Some(cfg.file.clone())
        };
        let s = Stats {
            enabled: cfg.enabled,
            file,
            ring_size: cfg.ring_size.max(1),
            retention_secs: parse_duration_secs(&cfg.retention).unwrap_or(7 * 86_400),
            max_file_bytes: cfg.max_file_mb.saturating_mul(1024 * 1024),
            inner: Mutex::new(Inner::default()),
        };
        s.load_folded();
        s.replay();
        s.maybe_compact(true);
        std::sync::Arc::new(s)
    }

    /// Replay the JSONL file into aggregates (within the retention window only).
    fn replay(&self) {
        let Some(path) = &self.file else { return };
        let Ok(content) = std::fs::read_to_string(path) else {
            return;
        };
        let cutoff = now_secs().saturating_sub(self.retention_secs);
        let mut inner = self.inner.lock().unwrap();
        for line in content.lines() {
            let line = line.trim();
            if line.is_empty() {
                continue;
            }
            if let Ok(rec) = serde_json::from_str::<RequestRecord>(line) {
                // push everything to the "recent" ring (it self-limits by size),
                // but apply to aggregates/series only entries within the window.
                Self::push_recent(&mut inner, &rec, self.ring_size);
                inner
                    .by_account_alltime
                    .entry(rec.account.clone())
                    .or_default()
                    .apply(&rec);
                Self::apply_month(&mut inner, &rec);
                if rec.ts >= cutoff {
                    Self::apply_aggregates(&mut inner, &rec, self.retention_secs);
                }
            }
        }
    }

    fn apply_month(inner: &mut Inner, rec: &RequestRecord) {
        if rec.ts <= inner.month_upto_ts || rec.cost_usd <= 0.0 {
            return;
        }
        let m = month_of(rec.ts);
        let e = inner
            .month_spend
            .entry(rec.account.clone())
            .or_insert((m.clone(), 0.0));
        if e.0 != m {
            *e = (m, 0.0);
        }
        e.1 += rec.cost_usd;
    }

    /// Month-to-date spend for an account (budget checks).
    pub fn month_spend(&self, account: &str) -> f64 {
        let inner = self.inner.lock().unwrap();
        let m = month_of(now_secs());
        inner
            .month_spend
            .get(account)
            .filter(|(mm, _)| *mm == m)
            .map(|(_, v)| *v)
            .unwrap_or(0.0)
    }

    fn push_recent(inner: &mut Inner, rec: &RequestRecord, ring_size: usize) {
        inner.recent.push_back(rec.clone());
        while inner.recent.len() > ring_size {
            inner.recent.pop_front();
        }
    }

    fn apply_aggregates(inner: &mut Inner, rec: &RequestRecord, retention_secs: u64) {
        // all-time counters for Prometheus; the windowed snapshot is derived
        // from the buckets below.
        inner.total.apply(rec);
        inner
            .by_model
            .entry(rec.model_id.clone())
            .or_default()
            .apply(rec);

        let minute = rec.ts - (rec.ts % 60);
        match inner.buckets.back_mut() {
            Some(b) if b.minute == minute => fill_bucket(b, rec),
            Some(b) if b.minute > minute => { /* arrived older than the tail — ignore for series */
            }
            _ => {
                let mut b = Bucket {
                    minute,
                    ..Default::default()
                };
                fill_bucket(&mut b, rec);
                inner.buckets.push_back(b);
            }
        }
        // trim old buckets outside the retention window
        let cutoff = now_secs().saturating_sub(retention_secs);
        while let Some(front) = inner.buckets.front() {
            if front.minute < cutoff {
                inner.buckets.pop_front();
            } else {
                break;
            }
        }
    }

    /// Record a new request event.
    pub fn record(&self, rec: RequestRecord) {
        if !self.enabled {
            return;
        }
        {
            let mut inner = self.inner.lock().unwrap();
            Self::apply_aggregates(&mut inner, &rec, self.retention_secs);
            Self::push_recent(&mut inner, &rec, self.ring_size);
            inner
                .by_account_alltime
                .entry(rec.account.clone())
                .or_default()
                .apply(&rec);
            Self::apply_month(&mut inner, &rec);
            inner.writes_since_check += 1;
        }
        if let Some(path) = &self.file {
            if let Ok(line) = serde_json::to_string(&rec) {
                if let Ok(mut f) = std::fs::OpenOptions::new()
                    .create(true)
                    .append(true)
                    .open(path)
                {
                    let _ = writeln!(f, "{line}");
                }
            }
        }
        self.maybe_compact(false);
    }

    fn totals_path(&self) -> Option<String> {
        self.file.as_ref().map(|f| format!("{f}.totals.json"))
    }

    fn load_folded(&self) {
        let Some(path) = self.totals_path() else {
            return;
        };
        let Ok(content) = std::fs::read_to_string(&path) else {
            return;
        };
        #[derive(serde::Deserialize)]
        struct Sidecar {
            #[serde(default)]
            by_account: HashMap<String, Totals>,
            #[serde(default)]
            month_spend: HashMap<String, (String, f64)>,
            #[serde(default)]
            month_upto_ts: u64,
        }
        // new format {by_account, month_spend, month_upto_ts}; the original
        // file held the bare by_account map — accept both
        if let Ok(sc) = serde_json::from_str::<Sidecar>(&content) {
            let mut inner = self.inner.lock().unwrap();
            inner.folded = sc.by_account;
            inner.month_spend = sc.month_spend;
            inner.month_upto_ts = sc.month_upto_ts;
        } else if let Ok(m) = serde_json::from_str::<HashMap<String, Totals>>(&content) {
            self.inner.lock().unwrap().folded = m;
        }
    }

    fn write_sidecar(&self, inner: &Inner) {
        let Some(tpath) = self.totals_path() else {
            return;
        };
        let payload = serde_json::json!({
            "by_account": inner.folded,
            "month_spend": inner.month_spend,
            "month_upto_ts": now_secs(),
        });
        let tmp = format!("{tpath}.tmp");
        if std::fs::write(&tmp, payload.to_string()).is_ok() {
            let _ = std::fs::rename(&tmp, &tpath);
        }
    }

    /// Keep the JSONL bounded: once it outgrows `max_file_bytes`, records
    /// older than the retention window fold into the cumulative per-account
    /// sidecar and are dropped from the file (tmp + rename on both files, so
    /// a crash mid-compaction loses nothing). `force` checks size regardless
    /// of the write counter (startup).
    fn maybe_compact(&self, force: bool) {
        if self.max_file_bytes == 0 {
            return;
        }
        let Some(path) = &self.file else { return };
        {
            let mut inner = self.inner.lock().unwrap();
            if !force {
                if inner.writes_since_check < 256 {
                    return;
                }
                inner.writes_since_check = 0;
                // piggyback: persist month-to-date spend so budgets survive
                // restarts even between compactions
                inner.month_upto_ts = now_secs();
                self.write_sidecar(&inner);
            }
        }
        let size = std::fs::metadata(path).map(|m| m.len()).unwrap_or(0);
        if size <= self.max_file_bytes {
            return;
        }
        let Ok(content) = std::fs::read_to_string(path) else {
            return;
        };
        let cutoff = now_secs().saturating_sub(self.retention_secs);
        let mut kept = String::with_capacity(content.len() / 2);
        let mut inner = self.inner.lock().unwrap();
        for line in content.lines() {
            let line = line.trim();
            if line.is_empty() {
                continue;
            }
            match serde_json::from_str::<RequestRecord>(line) {
                Ok(rec) if rec.ts < cutoff => {
                    inner
                        .folded
                        .entry(rec.account.clone())
                        .or_default()
                        .apply(&rec);
                    // fold = leave by_account_alltime (record leaves the file,
                    // its share moves from "alltime" into "folded")
                    if let Some(t) = inner.by_account_alltime.get_mut(&rec.account) {
                        // subtract by rebuilding is costly; instead track the
                        // invariant additively: remove the record's share
                        t.requests = t.requests.saturating_sub(1);
                        if rec.outcome == "ok" {
                            t.ok = t.ok.saturating_sub(1);
                        } else {
                            t.errors = t.errors.saturating_sub(1);
                        }
                        if rec.failover {
                            t.failovers = t.failovers.saturating_sub(1);
                        }
                        t.prompt_tokens = t.prompt_tokens.saturating_sub(rec.prompt_tokens as u64);
                        t.completion_tokens = t
                            .completion_tokens
                            .saturating_sub(rec.completion_tokens as u64);
                        t.cost_usd = (t.cost_usd - rec.cost_usd).max(0.0);
                        t.latency_ms_sum = t.latency_ms_sum.saturating_sub(rec.latency_ms);
                    }
                }
                Ok(_) => {
                    kept.push_str(line);
                    kept.push('\n');
                }
                Err(_) => { /* unparseable line — drop */ }
            }
        }
        // persist the folded totals FIRST, then swap the trimmed file in
        inner.month_upto_ts = now_secs();
        self.write_sidecar(&inner);
        let tmp = format!("{path}.tmp");
        if std::fs::write(&tmp, kept).is_ok() {
            let _ = std::fs::rename(&tmp, path);
        }
        tracing::info!(
            new_size = std::fs::metadata(path).map(|m| m.len()).unwrap_or(0),
            "stats log compacted (older records folded into cumulative totals)"
        );
    }

    /// Snapshot for the admin API: aggregates + time series for `range_secs` + breakdown by `groupby`.
    pub fn snapshot(&self, range_secs: u64, groupby: &str) -> serde_json::Value {
        let inner = self.inner.lock().unwrap();
        let cutoff = now_secs().saturating_sub(range_secs);
        let series: Vec<&Bucket> = inner
            .buckets
            .iter()
            .filter(|b| b.minute >= cutoff)
            .collect();

        // Windowed totals + breakdown = sum of the in-window buckets, so the KPIs
        // and the breakdown recompute whenever the range changes.
        let mut win = Totals::default();
        let mut group: HashMap<String, Totals> = HashMap::new();
        for b in &series {
            win.requests += b.requests;
            win.errors += b.errors;
            win.ok += b.requests.saturating_sub(b.errors);
            win.failovers += b.failovers;
            win.prompt_tokens += b.prompt_tokens;
            win.completion_tokens += b.completion_tokens;
            win.cost_usd += b.cost_usd;
            win.latency_ms_sum += b.latency_ms_sum;
            let src = match groupby {
                "account" => &b.by_account,
                "tier" => &b.by_tier,
                "label" => &b.by_label,
                _ => &b.by_model,
            };
            for (k, t) in src {
                group.entry(k.clone()).or_default().merge(t);
            }
        }

        let mut breakdown: Vec<serde_json::Value> = group
            .iter()
            .map(|(k, t)| {
                serde_json::json!({
                    "key": k,
                    "requests": t.requests,
                    "ok": t.ok,
                    "errors": t.errors,
                    "prompt_tokens": t.prompt_tokens,
                    "completion_tokens": t.completion_tokens,
                    "cost_usd": t.cost_usd,
                    "avg_latency_ms": avg(t.latency_ms_sum, t.requests),
                })
            })
            .collect();
        breakdown.sort_by(|a, b| {
            b["requests"]
                .as_u64()
                .unwrap_or(0)
                .cmp(&a["requests"].as_u64().unwrap_or(0))
        });

        let mut alltime: HashMap<String, Totals> = inner.folded.clone();
        for (k, t) in &inner.by_account_alltime {
            alltime.entry(k.clone()).or_default().merge(t);
        }
        let alltime_by_account: serde_json::Map<String, serde_json::Value> = alltime
            .iter()
            .map(|(k, t)| {
                (
                    k.clone(),
                    serde_json::json!({
                        "requests": t.requests,
                        "prompt_tokens": t.prompt_tokens,
                        "completion_tokens": t.completion_tokens,
                        "cost_usd": t.cost_usd,
                    }),
                )
            })
            .collect();

        serde_json::json!({
            "alltime_by_account": alltime_by_account,
            "totals": {
                "requests": win.requests,
                "ok": win.ok,
                "errors": win.errors,
                "failovers": win.failovers,
                "prompt_tokens": win.prompt_tokens,
                "completion_tokens": win.completion_tokens,
                "total_tokens": win.prompt_tokens + win.completion_tokens,
                "cost_usd": win.cost_usd,
                "avg_latency_ms": avg(win.latency_ms_sum, win.requests),
                "success_rate": if win.requests > 0 {
                    win.ok as f64 / win.requests as f64
                } else { 0.0 },
            },
            "groupby": groupby,
            "breakdown": breakdown,
            "series": series,
        })
    }

    /// Reset all in-memory stats and truncate the JSONL log — the "clear logs"
    /// admin action.
    pub fn clear(&self) {
        if let Some(tpath) = self.totals_path() {
            let _ = std::fs::remove_file(tpath);
        }
        {
            let mut inner = self.inner.lock().unwrap();
            *inner = Inner::default();
        }
        if let Some(path) = &self.file {
            let _ = std::fs::write(path, "");
        }
    }

    /// Recent requests (newest first), with pagination.
    pub fn recent(&self, limit: usize, offset: usize) -> Vec<RequestRecord> {
        let inner = self.inner.lock().unwrap();
        inner
            .recent
            .iter()
            .rev()
            .skip(offset)
            .take(limit)
            .cloned()
            .collect()
    }

    /// Metrics text in Prometheus format.
    pub fn prometheus(&self) -> String {
        let inner = self.inner.lock().unwrap();
        let mut out = String::new();
        let t = &inner.total;
        out.push_str("# HELP gw_requests_total Total gateway requests.\n");
        out.push_str("# TYPE gw_requests_total counter\n");
        out.push_str(&format!("gw_requests_total {}\n", t.requests));
        out.push_str("# HELP gw_failovers_total Total provider failovers.\n");
        out.push_str("# TYPE gw_failovers_total counter\n");
        out.push_str(&format!("gw_failovers_total {}\n", t.failovers));
        out.push_str("# HELP gw_tokens_total Total tokens by direction.\n");
        out.push_str("# TYPE gw_tokens_total counter\n");
        out.push_str(&format!(
            "gw_tokens_total{{direction=\"in\"}} {}\n",
            t.prompt_tokens
        ));
        out.push_str(&format!(
            "gw_tokens_total{{direction=\"out\"}} {}\n",
            t.completion_tokens
        ));
        out.push_str("# HELP gw_cost_usd_total Total estimated cost in USD.\n");
        out.push_str("# TYPE gw_cost_usd_total counter\n");
        out.push_str(&format!("gw_cost_usd_total {:.6}\n", t.cost_usd));
        out.push_str("# HELP gw_provider_calls_total Calls by model and outcome.\n");
        out.push_str("# TYPE gw_provider_calls_total counter\n");
        for (model, mt) in &inner.by_model {
            let model = escape_label(model);
            out.push_str(&format!(
                "gw_provider_calls_total{{model_id=\"{model}\",outcome=\"ok\"}} {}\n",
                mt.ok
            ));
            out.push_str(&format!(
                "gw_provider_calls_total{{model_id=\"{model}\",outcome=\"error\"}} {}\n",
                mt.errors
            ));
        }
        out
    }
}

fn fill_bucket(b: &mut Bucket, r: &RequestRecord) {
    b.requests += 1;
    if r.outcome != "ok" {
        b.errors += 1;
    }
    if r.failover {
        b.failovers += 1;
    }
    b.prompt_tokens += r.prompt_tokens as u64;
    b.completion_tokens += r.completion_tokens as u64;
    b.cost_usd += r.cost_usd;
    b.latency_ms_sum += r.latency_ms;
    // per-group breakdown, so the snapshot can be windowed by summing buckets
    b.by_model.entry(r.model_id.clone()).or_default().apply(r);
    b.by_account
        .entry(if r.account.is_empty() {
            "anonymous".to_string()
        } else {
            r.account.clone()
        })
        .or_default()
        .apply(r);
    b.by_tier.entry(r.tier.clone()).or_default().apply(r);
    b.by_label.entry(r.task_label.clone()).or_default().apply(r);
}

fn month_of(ts: u64) -> String {
    // календарный месяц UTC без внешних зависимостей: дни с эпохи → (год, месяц)
    let days = ts / 86_400;
    let mut year = 1970u64;
    let mut d = days;
    loop {
        let leap = (year % 4 == 0 && year % 100 != 0) || year % 400 == 0;
        let len = if leap { 366 } else { 365 };
        if d < len {
            break;
        }
        d -= len;
        year += 1;
    }
    let leap = (year % 4 == 0 && year % 100 != 0) || year % 400 == 0;
    let ml = [
        31,
        if leap { 29 } else { 28 },
        31,
        30,
        31,
        30,
        31,
        31,
        30,
        31,
        30,
        31,
    ];
    let mut month = 1;
    for len in ml {
        if d < len {
            break;
        }
        d -= len;
        month += 1;
    }
    format!("{year:04}-{month:02}")
}

fn avg(sum: u64, n: u64) -> u64 {
    sum.checked_div(n).unwrap_or(0)
}

fn escape_label(s: &str) -> String {
    s.replace('\\', "\\\\")
        .replace('"', "\\\"")
        .replace('\n', " ")
}

/// Parse a duration string of the form `7d` / `24h` / `90m` / `3600s` into seconds.
pub fn parse_duration_secs(s: &str) -> Option<u64> {
    let s = s.trim();
    if s.is_empty() {
        return None;
    }
    let (num, mult) = if let Some(n) = s.strip_suffix('d') {
        (n, 86_400)
    } else if let Some(n) = s.strip_suffix('h') {
        (n, 3_600)
    } else if let Some(n) = s.strip_suffix('m') {
        (n, 60)
    } else if let Some(n) = s.strip_suffix('s') {
        (n, 1)
    } else {
        (s, 1)
    };
    num.trim().parse::<u64>().ok().map(|v| v * mult)
}