remem/db/query/

stats.rs

use anyhow::Result;
use rusqlite::{params, Connection, OptionalExtension};

use crate::db::{
    AiUsageBreakdown, AiUsageSourceTotals, AiUsageTotals, DailyAiUsage, FailureLifecycleStats,
    WeeklyAiUsage,
};

pub use super::legacy_surfaces::LegacySurfaceStats;
use super::shared::collect_rows;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SystemStats {
    pub active_memories: i64,
    pub active_observations: i64,
    pub total_observations: i64,
    pub session_summaries: i64,
    pub raw_messages: i64,
    pub raw_ingest_failures: i64,
    pub raw_ingest_parse_errors: i64,
    pub raw_ingest_insert_errors: i64,
    pub latest_raw_ingest_failure_epoch: Option<i64>,
    pub latest_raw_ingest_failure_kind: Option<String>,
    pub latest_raw_ingest_failure_path: Option<String>,
    pub latest_raw_ingest_failure_message: Option<String>,
    pub captured_events: i64,
    pub latest_captured_event_epoch: Option<i64>,
    pub latest_capture_activity_epoch: Option<i64>,
    pub capture_drop_events: i64,
    pub actionable_capture_drops: i64,
    pub unrecovered_capture_spills: i64,
    pub latest_capture_drop_epoch: Option<i64>,
    pub latest_capture_drop_reason: Option<String>,
    pub latest_capture_drop_detail: Option<String>,
    pub pending_extraction_tasks: i64,
    pub processing_extraction_tasks: i64,
    pub expired_processing_extraction_tasks: i64,
    pub failed_extraction_tasks: i64,
    pub retryable_extraction_replay_ranges: i64,
    pub active_extraction_replay_ranges: i64,
    pub quarantined_extraction_replay_ranges: i64,
    pub oldest_pending_extraction_epoch: Option<i64>,
    pub pending_memory_candidates: i64,
    pub total_memory_candidates: i64,
    pub promoted_memory_candidates: i64,
    pub pending_review_memory_candidates: i64,
    pub pending_graph_candidates: i64,
    pub pending_observations: i64,
    pub ready_pending_observations: i64,
    pub delayed_pending_observations: i64,
    pub processing_pending_observations: i64,
    pub expired_processing_pending_observations: i64,
    pub failed_pending_observations: i64,
    pub oldest_ready_pending_epoch: Option<i64>,
    pub pending_jobs: i64,
    pub processing_jobs: i64,
    pub failed_jobs: i64,
    pub stuck_jobs: i64,
    pub failure_lifecycle: FailureLifecycleStats,
    pub worker_daemon_healthy: bool,
    pub worker_heartbeat_owner: Option<String>,
    pub worker_heartbeat_age_secs: Option<i64>,
    pub legacy_surfaces: Vec<LegacySurfaceStats>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MemoryFactsStats {
    pub table_exists: bool,
    pub total: i64,
    pub retrieval_eligible: i64,
    pub active_memories: i64,
    pub captured_events: i64,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DailyActivityStats {
    pub memories: i64,
    pub observations: i64,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProjectCount {
    pub project: String,
    pub count: i64,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CandidatePromotionStat {
    pub source_kind: String,
    pub review_status: String,
    pub block_reason: Option<String>,
    pub total: i64,
    pub last_7_days: i64,
}

pub fn query_system_stats(conn: &Connection) -> Result<SystemStats> {
    let now = chrono::Utc::now().timestamp();
    let raw_ingest = query_raw_ingest_failure_stats(conn)?;
    let capture_drop = crate::db::query_capture_drop_stats(conn)?;
    let captured_events =
        conn.query_row("SELECT COUNT(*) FROM captured_events", [], |row| row.get(0))?;
    let latest_captured_event_epoch = conn.query_row(
        "SELECT MAX(inserted_at_epoch) FROM captured_events",
        [],
        |row| row.get(0),
    )?;
    let latest_raw_message_epoch = conn.query_row(
        "SELECT MAX(created_at_epoch) FROM raw_messages",
        [],
        |row| row.get::<_, Option<i64>>(0),
    )?;
    let latest_capture_activity_epoch = [
        latest_captured_event_epoch,
        capture_drop.latest_epoch,
        latest_raw_message_epoch,
    ]
    .into_iter()
    .flatten()
    .max();
    let replay_ranges = query_extraction_replay_range_stats(conn)?;
    let failure_lifecycle = crate::db::query_failure_lifecycle_stats(conn, now)?;
    let worker_heartbeat = crate::db::worker::latest_daemon_worker_heartbeat(conn)?;
    let healthy_worker_heartbeat = crate::db::worker::healthy_daemon_worker_heartbeat(
        conn,
        crate::db::worker::WORKER_HEARTBEAT_HEALTH_SECS,
    )?;
    let worker_heartbeat_age_secs = worker_heartbeat
        .as_ref()
        .map(|heartbeat| now.saturating_sub(heartbeat.updated_at_epoch));
    let worker_daemon_healthy = healthy_worker_heartbeat.is_some();
    let legacy_surfaces = super::legacy_surfaces::query_legacy_surface_stats(conn)?;
    Ok(SystemStats {
        active_memories: query_current_active_memory_count(conn)?,
        active_observations: conn.query_row(
            "SELECT COUNT(*) FROM observations WHERE status = 'active'",
            [],
            |row| row.get(0),
        )?,
        total_observations: conn.query_row("SELECT COUNT(*) FROM observations", [], |row| {
            row.get(0)
        })?,
        session_summaries: conn.query_row("SELECT COUNT(*) FROM session_summaries", [], |row| {
            row.get(0)
        })?,
        raw_messages: conn.query_row("SELECT COUNT(*) FROM raw_messages", [], |row| row.get(0))?,
        raw_ingest_failures: raw_ingest.failures,
        raw_ingest_parse_errors: raw_ingest.parse_errors,
        raw_ingest_insert_errors: raw_ingest.insert_errors,
        latest_raw_ingest_failure_epoch: raw_ingest.latest_epoch,
        latest_raw_ingest_failure_kind: raw_ingest.latest_kind,
        latest_raw_ingest_failure_path: raw_ingest.latest_path,
        latest_raw_ingest_failure_message: raw_ingest.latest_message,
        captured_events,
        latest_captured_event_epoch,
        latest_capture_activity_epoch,
        capture_drop_events: capture_drop.total,
        actionable_capture_drops: capture_drop.actionable,
        unrecovered_capture_spills: capture_drop.unrecovered_spills,
        latest_capture_drop_epoch: capture_drop.latest_epoch,
        latest_capture_drop_reason: capture_drop.latest_reason,
        latest_capture_drop_detail: capture_drop.latest_detail,
        pending_extraction_tasks: conn.query_row(
            "SELECT COUNT(*) FROM extraction_tasks WHERE status = 'pending'",
            [],
            |row| row.get(0),
        )?,
        processing_extraction_tasks: conn.query_row(
            "SELECT COUNT(*) FROM extraction_tasks WHERE status = 'processing'",
            [],
            |row| row.get(0),
        )?,
        expired_processing_extraction_tasks: conn.query_row(
            "SELECT COUNT(*)
             FROM extraction_tasks
             WHERE status = 'processing'
               AND lease_expires_epoch IS NOT NULL
               AND lease_expires_epoch < ?1",
            params![now],
            |row| row.get(0),
        )?,
        failed_extraction_tasks: failure_lifecycle.extraction_task.actionable_total,
        retryable_extraction_replay_ranges: replay_ranges.retryable,
        active_extraction_replay_ranges: replay_ranges.active,
        quarantined_extraction_replay_ranges: replay_ranges.quarantined,
        oldest_pending_extraction_epoch: conn.query_row(
            "SELECT MIN(created_at_epoch) FROM extraction_tasks WHERE status = 'pending'",
            [],
            |row| row.get(0),
        )?,
        pending_memory_candidates: conn.query_row(
            "SELECT COUNT(*) FROM memory_candidates WHERE review_status = 'pending_review'",
            [],
            |row| row.get(0),
        )?,
        total_memory_candidates: conn.query_row("SELECT COUNT(*) FROM memory_candidates", [], |row| {
            row.get(0)
        })?,
        promoted_memory_candidates: conn.query_row(
            "SELECT COUNT(*) FROM memory_candidates
             WHERE review_status IN ('auto_promoted', 'approved', 'edited')",
            [],
            |row| row.get(0),
        )?,
        pending_review_memory_candidates: conn.query_row(
            "SELECT COUNT(*) FROM memory_candidates WHERE review_status = 'pending_review'",
            [],
            |row| row.get(0),
        )?,
        pending_graph_candidates: query_pending_graph_candidates(conn)?,
        pending_observations: conn.query_row(
            "SELECT COUNT(*) FROM pending_observations WHERE status = 'pending'",
            [],
            |row| row.get(0),
        )?,
        ready_pending_observations: conn.query_row(
            "SELECT COUNT(*) FROM pending_observations
             WHERE status = 'pending'
               AND (next_retry_epoch IS NULL OR next_retry_epoch <= ?1)
               AND (lease_owner IS NULL OR lease_expires_epoch IS NULL OR lease_expires_epoch < ?1)",
            params![now],
            |row| row.get(0),
        )?,
        delayed_pending_observations: conn.query_row(
            "SELECT COUNT(*) FROM pending_observations
             WHERE status = 'pending'
               AND next_retry_epoch IS NOT NULL
               AND next_retry_epoch > ?1",
            params![now],
            |row| row.get(0),
        )?,
        processing_pending_observations: conn.query_row(
            "SELECT COUNT(*) FROM pending_observations WHERE status = 'processing'",
            [],
            |row| row.get(0),
        )?,
        expired_processing_pending_observations: conn.query_row(
            "SELECT COUNT(*) FROM pending_observations
             WHERE status = 'processing'
               AND lease_expires_epoch IS NOT NULL
               AND lease_expires_epoch < ?1",
            params![now],
            |row| row.get(0),
        )?,
        failed_pending_observations: failure_lifecycle.pending_observation.actionable_total,
        oldest_ready_pending_epoch: conn.query_row(
            "SELECT MIN(created_at_epoch) FROM pending_observations
             WHERE status = 'pending'
               AND (next_retry_epoch IS NULL OR next_retry_epoch <= ?1)
               AND (lease_owner IS NULL OR lease_expires_epoch IS NULL OR lease_expires_epoch < ?1)",
            params![now],
            |row| row.get(0),
        )?,
        pending_jobs: conn.query_row("SELECT COUNT(*) FROM jobs WHERE state = 'pending'", [], |row| {
            row.get(0)
        })?,
        processing_jobs: conn.query_row(
            "SELECT COUNT(*) FROM jobs WHERE state = 'processing'",
            [],
            |row| row.get(0),
        )?,
        failed_jobs: failure_lifecycle.job.actionable_total,
        stuck_jobs: conn.query_row(
            "SELECT COUNT(*) FROM jobs WHERE state = 'processing' \
             AND lease_expires_epoch < strftime('%s', 'now')",
            [],
            |row| row.get(0),
        )?,
        worker_daemon_healthy,
        worker_heartbeat_owner: worker_heartbeat.map(|heartbeat| heartbeat.owner),
        worker_heartbeat_age_secs,
        failure_lifecycle,
        legacy_surfaces,
    })
}

pub fn query_memory_facts_stats(conn: &Connection) -> Result<MemoryFactsStats> {
    let memory_facts_exists = table_exists(conn, "memory_facts")?;
    let memories_exists = table_exists(conn, "memories")?;
    let captured_events_exists = table_exists(conn, "captured_events")?;
    let total = if memory_facts_exists {
        conn.query_row("SELECT COUNT(*) FROM memory_facts", [], |row| row.get(0))?
    } else {
        0
    };
    let retrieval_eligible = if memory_facts_exists && memories_exists {
        let fact_current_filter = crate::memory::facts::current_fact_filter_sql(
            "f",
            crate::memory::facts::invalidated_at_epoch_available(conn)?,
        );
        let sql = format!(
            "SELECT COUNT(*)
             FROM memory_facts f
             JOIN memories m ON m.id = f.source_memory_id
             WHERE {fact_current_filter}
               AND f.valid_from_epoch IS NOT NULL
               AND m.status = 'active'
               AND (
                 m.expires_at_epoch IS NULL
                 OR m.expires_at_epoch > CAST(strftime('%s', 'now') AS INTEGER)
               )"
        );
        conn.query_row(&sql, [], |row| row.get(0))?
    } else {
        0
    };
    let active_memories = if memories_exists {
        query_current_active_memory_count(conn)?
    } else {
        0
    };
    let captured_events = if captured_events_exists {
        conn.query_row("SELECT COUNT(*) FROM captured_events", [], |row| row.get(0))?
    } else {
        0
    };
    Ok(MemoryFactsStats {
        table_exists: memory_facts_exists,
        total,
        retrieval_eligible,
        active_memories,
        captured_events,
    })
}

#[derive(Debug, Clone, Default)]
struct RawIngestFailureStats {
    failures: i64,
    parse_errors: i64,
    insert_errors: i64,
    latest_epoch: Option<i64>,
    latest_kind: Option<String>,
    latest_path: Option<String>,
    latest_message: Option<String>,
}

#[derive(Debug, Clone, Default)]
struct ExtractionReplayRangeStats {
    retryable: i64,
    active: i64,
    quarantined: i64,
}

fn query_extraction_replay_range_stats(conn: &Connection) -> Result<ExtractionReplayRangeStats> {
    if !table_exists(conn, "extraction_replay_ranges")? {
        return Ok(ExtractionReplayRangeStats::default());
    }
    let archived_filter = if column_exists(conn, "extraction_replay_ranges", "archived_at_epoch")? {
        "AND r.archived_at_epoch IS NULL"
    } else {
        ""
    };
    let archived_filter_bare =
        if column_exists(conn, "extraction_replay_ranges", "archived_at_epoch")? {
            "AND archived_at_epoch IS NULL"
        } else {
            ""
        };

    Ok(ExtractionReplayRangeStats {
        retryable: conn.query_row(
            &format!(
                "SELECT COUNT(*)
             FROM extraction_replay_ranges r
             WHERE r.status IN ('pending', 'failed')
               {archived_filter}
               AND NOT EXISTS (
                 SELECT 1
                 FROM extraction_tasks t
                 WHERE t.replay_range_id = r.id
                   AND t.status IN ('pending', 'processing')
               )"
            ),
            [],
            |row| row.get(0),
        )?,
        active: conn.query_row(
            &format!(
                "SELECT COUNT(*)
             FROM extraction_replay_ranges r
             WHERE r.status = 'requeued'
                {archived_filter}
                OR (r.status IN ('pending', 'failed') AND EXISTS (
                 SELECT 1
                 FROM extraction_tasks t
                 WHERE t.replay_range_id = r.id
                   AND t.status IN ('pending', 'processing')
               ) {archived_filter})"
            ),
            [],
            |row| row.get(0),
        )?,
        quarantined: conn.query_row(
            &format!(
                "SELECT COUNT(*)
             FROM extraction_replay_ranges
             WHERE status = 'quarantined'
               {archived_filter_bare}"
            ),
            [],
            |row| row.get(0),
        )?,
    })
}

fn query_raw_ingest_failure_stats(conn: &Connection) -> Result<RawIngestFailureStats> {
    if !table_exists(conn, "raw_ingest_failures")? {
        return Ok(RawIngestFailureStats::default());
    }

    let (failures, parse_errors, insert_errors) = conn.query_row(
        "SELECT COUNT(*), COALESCE(SUM(parse_errors), 0), COALESCE(SUM(insert_errors), 0)
         FROM raw_ingest_failures",
        [],
        |row| {
            Ok((
                row.get::<_, i64>(0)?,
                row.get::<_, i64>(1)?,
                row.get::<_, i64>(2)?,
            ))
        },
    )?;
    let latest = conn
        .query_row(
            "SELECT created_at_epoch, error_kind, transcript_path, error_message
             FROM raw_ingest_failures
             ORDER BY created_at_epoch DESC, id DESC
             LIMIT 1",
            [],
            |row| {
                Ok((
                    row.get::<_, i64>(0)?,
                    row.get::<_, String>(1)?,
                    row.get::<_, Option<String>>(2)?,
                    row.get::<_, String>(3)?,
                ))
            },
        )
        .optional()?;

    let (latest_epoch, latest_kind, latest_path, latest_message) = match latest {
        Some((epoch, kind, path, message)) => (Some(epoch), Some(kind), path, Some(message)),
        None => (None, None, None, None),
    };

    Ok(RawIngestFailureStats {
        failures,
        parse_errors,
        insert_errors,
        latest_epoch,
        latest_kind,
        latest_path,
        latest_message,
    })
}

fn query_current_active_memory_count(conn: &Connection) -> Result<i64> {
    Ok(conn.query_row(
        "SELECT COUNT(*)
         FROM memories
         WHERE status = 'active'
           AND (
             expires_at_epoch IS NULL
             OR expires_at_epoch > CAST(strftime('%s', 'now') AS INTEGER)
           )",
        [],
        |row| row.get(0),
    )?)
}

fn table_exists(conn: &Connection, table: &str) -> Result<bool> {
    Ok(conn
        .query_row(
            "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?1",
            [table],
            |_| Ok(()),
        )
        .optional()?
        .is_some())
}

fn column_exists(conn: &Connection, table: &str, column: &str) -> Result<bool> {
    if !table_exists(conn, table)? {
        return Ok(false);
    }
    let mut stmt = conn.prepare(&format!("PRAGMA table_info({})", quote_identifier(table)))?;
    let mut rows = stmt.query([])?;
    while let Some(row) = rows.next()? {
        let name: String = row.get(1)?;
        if name == column {
            return Ok(true);
        }
    }
    Ok(false)
}

fn quote_identifier(identifier: &str) -> String {
    format!("\"{}\"", identifier.replace('"', "\"\""))
}

fn query_pending_graph_candidates(conn: &Connection) -> Result<i64> {
    if !table_exists(conn, "graph_candidates")? {
        return Ok(0);
    }
    conn.query_row(
		"SELECT COUNT(*) FROM graph_candidates WHERE review_status IN ('pending_review', 'deferred')",
		[],
		|row| row.get(0),
	)
    .map_err(Into::into)
}

pub fn query_daily_activity_stats(
    conn: &Connection,
    since_epoch: i64,
) -> Result<DailyActivityStats> {
    Ok(DailyActivityStats {
        memories: conn.query_row(
            "SELECT COUNT(*) FROM memories WHERE created_at_epoch >= ?1",
            params![since_epoch],
            |row| row.get(0),
        )?,
        observations: conn.query_row(
            "SELECT COUNT(*) FROM observations WHERE created_at_epoch >= ?1",
            params![since_epoch],
            |row| row.get(0),
        )?,
    })
}

pub fn query_candidate_promotion_stats(
    conn: &Connection,
    now_epoch: i64,
) -> Result<Vec<CandidatePromotionStat>> {
    let week_ago = now_epoch - 7 * 24 * 3600;
    let mut stmt = conn.prepare(
        "SELECT source_kind,
                review_status,
                auto_promote_block_reason,
                COUNT(*) AS total,
                SUM(CASE WHEN created_at_epoch >= ?1 THEN 1 ELSE 0 END) AS last_7_days
         FROM memory_candidates
         GROUP BY source_kind, review_status, auto_promote_block_reason
         ORDER BY total DESC, source_kind ASC, review_status ASC, auto_promote_block_reason ASC",
    )?;
    let rows = stmt.query_map(params![week_ago], |row| {
        Ok(CandidatePromotionStat {
            source_kind: row.get(0)?,
            review_status: row.get(1)?,
            block_reason: row.get(2)?,
            total: row.get(3)?,
            last_7_days: row.get(4)?,
        })
    })?;
    collect_rows(rows)
}

pub fn query_top_projects(conn: &Connection, limit: i64) -> Result<Vec<ProjectCount>> {
    let mut stmt = conn.prepare(
        "SELECT project, COUNT(*) as cnt
         FROM memories
         WHERE status = 'active'
           AND (
             expires_at_epoch IS NULL
             OR expires_at_epoch > CAST(strftime('%s', 'now') AS INTEGER)
           )
         GROUP BY project ORDER BY cnt DESC, project ASC LIMIT ?1",
    )?;
    let rows = stmt.query_map(params![limit], |row| {
        Ok(ProjectCount {
            project: row.get(0)?,
            count: row.get(1)?,
        })
    })?;
    collect_rows(rows)
}

pub fn query_ai_usage_totals(
    conn: &Connection,
    since_epoch: Option<i64>,
    project: Option<&str>,
) -> Result<AiUsageTotals> {
    conn.query_row(
        "SELECT COUNT(*),
                COALESCE(SUM(input_tokens), 0),
                COALESCE(SUM(output_tokens), 0),
                COALESCE(SUM(reasoning_tokens), 0),
                COALESCE(SUM(cache_creation_tokens), 0),
                COALESCE(SUM(cache_read_tokens), 0),
                COALESCE(SUM(total_tokens), 0),
                COALESCE(SUM(estimated_cost_usd), 0.0)
         FROM ai_usage_events
         WHERE (?1 IS NULL OR created_at_epoch >= ?1)
           AND (?2 IS NULL OR project = ?2)",
        params![since_epoch, project],
        |row| {
            Ok(AiUsageTotals {
                calls: row.get(0)?,
                input_tokens: row.get(1)?,
                output_tokens: row.get(2)?,
                reasoning_tokens: row.get(3)?,
                cache_creation_tokens: row.get(4)?,
                cache_read_tokens: row.get(5)?,
                total_tokens: row.get(6)?,
                estimated_cost_usd: row.get(7)?,
            })
        },
    )
    .map_err(Into::into)
}

pub fn query_ai_usage_source_totals(
    conn: &Connection,
    since_epoch: Option<i64>,
    project: Option<&str>,
) -> Result<Vec<AiUsageSourceTotals>> {
    let mut stmt = conn.prepare(
        "SELECT usage_source,
                pricing_source,
                COUNT(*),
                COALESCE(SUM(total_tokens), 0),
                COALESCE(SUM(estimated_cost_usd), 0.0)
         FROM ai_usage_events
         WHERE (?1 IS NULL OR created_at_epoch >= ?1)
           AND (?2 IS NULL OR project = ?2)
         GROUP BY usage_source, pricing_source
         ORDER BY SUM(total_tokens) DESC",
    )?;
    let rows = stmt.query_map(params![since_epoch, project], |row| {
        Ok(AiUsageSourceTotals {
            usage_source: row.get(0)?,
            pricing_source: row.get(1)?,
            calls: row.get(2)?,
            total_tokens: row.get(3)?,
            estimated_cost_usd: row.get(4)?,
        })
    })?;
    collect_rows(rows)
}

pub fn query_ai_usage_breakdown(
    conn: &Connection,
    since_epoch: Option<i64>,
    project: Option<&str>,
    limit: i64,
) -> Result<Vec<AiUsageBreakdown>> {
    if limit <= 0 {
        return Ok(Vec::new());
    }

    let mut stmt = conn.prepare(
        "SELECT project,
                executor,
                usage_source,
                pricing_source,
                COUNT(*),
                COALESCE(SUM(total_tokens), 0),
                COALESCE(SUM(estimated_cost_usd), 0.0)
         FROM ai_usage_events
         WHERE (?1 IS NULL OR created_at_epoch >= ?1)
           AND (?2 IS NULL OR project = ?2)
         GROUP BY project, executor, usage_source, pricing_source
         ORDER BY SUM(estimated_cost_usd) DESC,
                  SUM(total_tokens) DESC,
                  COUNT(*) DESC,
                  project ASC,
                  executor ASC
         LIMIT ?3",
    )?;
    let rows = stmt.query_map(params![since_epoch, project, limit], |row| {
        Ok(AiUsageBreakdown {
            project: row.get(0)?,
            executor: row.get(1)?,
            usage_source: row.get(2)?,
            pricing_source: row.get(3)?,
            calls: row.get(4)?,
            total_tokens: row.get(5)?,
            estimated_cost_usd: row.get(6)?,
        })
    })?;
    collect_rows(rows)
}

pub fn query_daily_ai_usage(
    conn: &Connection,
    since_epoch: i64,
    project: Option<&str>,
    limit: i64,
) -> Result<Vec<DailyAiUsage>> {
    let mut stmt = conn.prepare(
        "SELECT strftime('%Y-%m-%d', created_at_epoch, 'unixepoch') AS day,
                COUNT(*),
                COALESCE(SUM(input_tokens), 0),
                COALESCE(SUM(output_tokens), 0),
                COALESCE(SUM(reasoning_tokens), 0),
                COALESCE(SUM(cache_creation_tokens), 0),
                COALESCE(SUM(cache_read_tokens), 0),
                COALESCE(SUM(total_tokens), 0),
                COALESCE(SUM(estimated_cost_usd), 0.0)
         FROM ai_usage_events
         WHERE created_at_epoch >= ?1
           AND (?2 IS NULL OR project = ?2)
         GROUP BY day
         ORDER BY day DESC
         LIMIT ?3",
    )?;
    let rows = stmt.query_map(params![since_epoch, project, limit], |row| {
        Ok(DailyAiUsage {
            day: row.get(0)?,
            calls: row.get(1)?,
            input_tokens: row.get(2)?,
            output_tokens: row.get(3)?,
            reasoning_tokens: row.get(4)?,
            cache_creation_tokens: row.get(5)?,
            cache_read_tokens: row.get(6)?,
            total_tokens: row.get(7)?,
            estimated_cost_usd: row.get(8)?,
        })
    })?;
    collect_rows(rows)
}

pub fn query_weekly_ai_usage(
    conn: &Connection,
    since_epoch: i64,
    project: Option<&str>,
    limit: i64,
) -> Result<Vec<WeeklyAiUsage>> {
    let mut stmt = conn.prepare(
        "SELECT strftime('%Y-W%W', created_at_epoch, 'unixepoch') AS week,
                COUNT(*),
                COALESCE(SUM(input_tokens), 0),
                COALESCE(SUM(output_tokens), 0),
                COALESCE(SUM(reasoning_tokens), 0),
                COALESCE(SUM(cache_creation_tokens), 0),
                COALESCE(SUM(cache_read_tokens), 0),
                COALESCE(SUM(total_tokens), 0),
                COALESCE(SUM(estimated_cost_usd), 0.0)
         FROM ai_usage_events
         WHERE created_at_epoch >= ?1
           AND (?2 IS NULL OR project = ?2)
         GROUP BY week
         ORDER BY week DESC
         LIMIT ?3",
    )?;
    let rows = stmt.query_map(params![since_epoch, project, limit], |row| {
        Ok(WeeklyAiUsage {
            week: row.get(0)?,
            calls: row.get(1)?,
            input_tokens: row.get(2)?,
            output_tokens: row.get(3)?,
            reasoning_tokens: row.get(4)?,
            cache_creation_tokens: row.get(5)?,
            cache_read_tokens: row.get(6)?,
            total_tokens: row.get(7)?,
            estimated_cost_usd: row.get(8)?,
        })
    })?;
    collect_rows(rows)
}

#[cfg(test)]
mod tests;