owlen/crates/owlen-core/src/usage.rs

use crate::{Error, Result};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, VecDeque};
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::fs;

const LEDGER_VERSION: u32 = 1;
const SECONDS_PER_HOUR: i64 = 60 * 60;
const SECONDS_PER_WEEK: i64 = 7 * 24 * 60 * 60;

#[derive(Clone, Debug, Serialize, Deserialize)]
struct UsageRecord {
    timestamp: i64,
    prompt_tokens: u32,
    completion_tokens: u32,
}

#[derive(Serialize, Deserialize)]
struct LedgerFile {
    version: u32,
    providers: HashMap<String, VecDeque<UsageRecord>>,
}

impl Default for LedgerFile {
    fn default() -> Self {
        Self {
            version: LEDGER_VERSION,
            providers: HashMap::new(),
        }
    }
}

#[derive(Clone, Debug, Default)]
pub struct UsageLedger {
    path: PathBuf,
    providers: HashMap<String, VecDeque<UsageRecord>>,
}

#[derive(Clone, Debug, Default)]
pub struct UsageQuota {
    pub hourly_quota_tokens: Option<u64>,
    pub weekly_quota_tokens: Option<u64>,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum UsageWindow {
    Hour,
    Week,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum UsageBand {
    Normal = 0,
    Warning = 1,
    Critical = 2,
}

#[derive(Clone, Debug, Default)]
pub struct WindowMetrics {
    pub prompt_tokens: u64,
    pub completion_tokens: u64,
    pub total_tokens: u64,
    pub quota_tokens: Option<u64>,
}

impl WindowMetrics {
    pub fn percent_of_quota(&self) -> Option<f64> {
        let quota = self.quota_tokens?;
        if quota == 0 {
            return None;
        }
        Some(self.total_tokens as f64 / quota as f64)
    }

    pub fn band(&self) -> UsageBand {
        match self.percent_of_quota() {
            Some(p) if p >= 0.95_f64 => UsageBand::Critical,
            Some(p) if p >= 0.80_f64 => UsageBand::Warning,
            _ => UsageBand::Normal,
        }
    }
}

#[derive(Clone, Debug, Default)]
pub struct UsageSnapshot {
    pub provider: String,
    pub hourly: WindowMetrics,
    pub weekly: WindowMetrics,
    pub last_updated: Option<SystemTime>,
}

impl UsageSnapshot {
    pub fn window(&self, window: UsageWindow) -> &WindowMetrics {
        match window {
            UsageWindow::Hour => &self.hourly,
            UsageWindow::Week => &self.weekly,
        }
    }
}

impl UsageLedger {
    pub fn empty(path: PathBuf) -> Self {
        Self {
            path,
            providers: HashMap::new(),
        }
    }

    pub async fn load_or_default(path: PathBuf) -> Result<Self> {
        if !path.exists() {
            return Ok(Self {
                path,
                providers: HashMap::new(),
            });
        }

        let contents = fs::read_to_string(&path)
            .await
            .map_err(|err| Error::Storage(format!("Failed to read usage ledger: {err}")))?;

        let file: LedgerFile = match serde_json::from_str(&contents) {
            Ok(file) => file,
            Err(err) => {
                return Err(Error::Storage(format!(
                    "Failed to parse usage ledger at {}: {err}",
                    path.display()
                )));
            }
        };

        Ok(Self {
            path,
            providers: file.providers,
        })
    }

    pub async fn persist(&self) -> Result<()> {
        if let Some(parent) = self.path.parent() {
            fs::create_dir_all(parent)
                .await
                .map_err(|err| Error::Storage(format!("Failed to create data directory: {err}")))?;
        }

        let serialized = self.serialize()?;

        fs::write(&self.path, serialized)
            .await
            .map_err(|err| Error::Storage(format!("Failed to write usage ledger: {err}")))?;

        Ok(())
    }

    pub fn record(
        &mut self,
        provider: &str,
        usage: &crate::types::TokenUsage,
        timestamp: SystemTime,
    ) {
        let total_tokens = usage.total_tokens;
        if total_tokens == 0 {
            return;
        }

        let ts = match timestamp.duration_since(UNIX_EPOCH) {
            Ok(duration) => duration.as_secs() as i64,
            Err(_) => 0,
        };

        let entry = self
            .providers
            .entry(provider.to_string())
            .or_insert_with(VecDeque::new);

        entry.push_back(UsageRecord {
            timestamp: ts,
            prompt_tokens: usage.prompt_tokens,
            completion_tokens: usage.completion_tokens,
        });

        self.prune_old(provider, ts);
    }

    pub fn provider_keys(&self) -> impl Iterator<Item = &String> {
        self.providers.keys()
    }

    pub fn serialize(&self) -> Result<String> {
        let file = LedgerFile {
            version: LEDGER_VERSION,
            providers: self.providers.clone(),
        };

        serde_json::to_string_pretty(&file)
            .map_err(|err| Error::Storage(format!("Failed to serialize usage ledger: {err}")))
    }

    pub fn path(&self) -> &Path {
        &self.path
    }

    pub fn snapshot(&self, provider: &str, quotas: UsageQuota, now: SystemTime) -> UsageSnapshot {
        let now_secs = now
            .duration_since(UNIX_EPOCH)
            .unwrap_or_else(|_| Duration::from_secs(0))
            .as_secs() as i64;

        let mut snapshot = UsageSnapshot {
            provider: provider.to_string(),
            hourly: WindowMetrics {
                quota_tokens: quotas.hourly_quota_tokens,
                ..Default::default()
            },
            weekly: WindowMetrics {
                quota_tokens: quotas.weekly_quota_tokens,
                ..Default::default()
            },
            last_updated: None,
        };

        if let Some(records) = self.providers.get(provider) {
            for record in records {
                if now_secs - record.timestamp <= SECONDS_PER_HOUR {
                    snapshot.hourly.prompt_tokens += record.prompt_tokens as u64;
                    snapshot.hourly.completion_tokens += record.completion_tokens as u64;
                }

                if now_secs - record.timestamp <= SECONDS_PER_WEEK {
                    snapshot.weekly.prompt_tokens += record.prompt_tokens as u64;
                    snapshot.weekly.completion_tokens += record.completion_tokens as u64;
                }
            }

            snapshot.hourly.total_tokens =
                snapshot.hourly.prompt_tokens + snapshot.hourly.completion_tokens;
            snapshot.weekly.total_tokens =
                snapshot.weekly.prompt_tokens + snapshot.weekly.completion_tokens;

            snapshot.last_updated = records.back().and_then(|record| {
                UNIX_EPOCH.checked_add(Duration::from_secs(record.timestamp as u64))
            });
        }

        snapshot
    }

    pub fn prune_old(&mut self, provider: &str, now_secs: i64) {
        if let Some(records) = self.providers.get_mut(provider) {
            while let Some(front) = records.front() {
                if now_secs - front.timestamp > SECONDS_PER_WEEK {
                    records.pop_front();
                } else {
                    break;
                }
            }
        }
    }

    pub fn prune_all(&mut self, now: SystemTime) {
        let now_secs = now
            .duration_since(UNIX_EPOCH)
            .unwrap_or_else(|_| Duration::from_secs(0))
            .as_secs() as i64;
        let provider_keys: Vec<String> = self.providers.keys().cloned().collect();
        for provider in provider_keys {
            self.prune_old(&provider, now_secs);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::TokenUsage;
    use std::time::{Duration, UNIX_EPOCH};
    use tempfile::tempdir;

    fn make_usage(prompt: u32, completion: u32) -> TokenUsage {
        TokenUsage {
            prompt_tokens: prompt,
            completion_tokens: completion,
            total_tokens: prompt.saturating_add(completion),
        }
    }

    #[test]
    fn records_and_summarizes_usage() {
        let temp = tempdir().expect("tempdir");
        let path = temp.path().join("ledger.json");
        let mut ledger = UsageLedger::empty(path);

        let usage = make_usage(40, 10);
        let timestamp = UNIX_EPOCH + Duration::from_secs(1);
        ledger.record("ollama_cloud", &usage, timestamp);

        let quotas = UsageQuota {
            hourly_quota_tokens: Some(100),
            weekly_quota_tokens: Some(1000),
        };

        let snapshot = ledger.snapshot("ollama_cloud", quotas, UNIX_EPOCH + Duration::from_secs(2));

        assert_eq!(snapshot.hourly.total_tokens, 50);
        assert_eq!(snapshot.weekly.total_tokens, 50);
        assert_eq!(snapshot.hourly.quota_tokens, Some(100));
        assert_eq!(snapshot.weekly.quota_tokens, Some(1000));
        assert_eq!(snapshot.hourly.band(), UsageBand::Normal);
    }

    #[test]
    fn prunes_records_outside_week() {
        let temp = tempdir().expect("tempdir");
        let path = temp.path().join("ledger.json");
        let mut ledger = UsageLedger::empty(path);

        let old_usage = make_usage(30, 5);
        let recent_usage = make_usage(20, 5);

        let base = UNIX_EPOCH;
        ledger.record("ollama_cloud", &old_usage, base);

        // Advance beyond a week for the second record.
        let later = UNIX_EPOCH + Duration::from_secs(SECONDS_PER_WEEK as u64 + 120);
        ledger.record("ollama_cloud", &recent_usage, later);

        let quotas = UsageQuota::default();
        let snapshot = ledger.snapshot("ollama_cloud", quotas, later);

        assert_eq!(snapshot.hourly.total_tokens, 25);
        assert_eq!(snapshot.weekly.total_tokens, 25);
    }
}