wsb-daemon 0.258.6441

use wsb::prelude::*;
use axum::{
    extract::{Query, State},
    http::StatusCode,
    Json,
};

use super::sessions::ErrorResponse;
use wsb::persist::global::models::{
    DailyUsage, ModelUsage, ProjectUsage, UsageFilter, UsageSummary,
};
use crate::web::machine_web::{place_filter, MachineWeb};
use wsb::persist::global::models::PlaceFilter;

/// The parameters the analytics handlers act on: what to count, and which
/// sessions by where they ran.
#[derive(Deserialize)]
pub struct AnalyticsQuery {
    pub include_one_shot: Option<String>,
    pub granularity: Option<String>,
    pub metric: Option<String>,
    pub project: Option<String>,
    pub worktree: Option<String>,
    pub placed: Option<String>,
}

impl AnalyticsQuery {
    fn place(&self) -> Result<PlaceFilter, (StatusCode, Json<ErrorResponse>)> {
        place_filter(self.project.clone(), self.worktree.clone(), self.placed.clone())
    }
}

/// `sql` with the place's conditions added where it says `{place}`, and the
/// values they bind.
fn placed(sql: &str, place: &PlaceFilter) -> (String, Vec<Box<dyn rusqlite::ToSql>>) {
    let mut conditions = String::new();
    let mut binds = Vec::new();
    place.push_sql("sessions", &mut conditions, &mut binds);
    (sql.replace("{place}", &conditions), binds)
}

fn bound(binds: &[Box<dyn rusqlite::ToSql>]) -> Vec<&dyn rusqlite::ToSql> {
    binds.iter().map(|b| b.as_ref()).collect()
}

fn internal_error(msg: &str) -> (StatusCode, Json<ErrorResponse>) {
    (
        StatusCode::INTERNAL_SERVER_ERROR,
        Json(ErrorResponse {
            error: msg.to_string(),
        }),
    )
}

pub async fn summary(
    State(state): State<MachineWeb>,
    Query(params): Query<AnalyticsQuery>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    let include_one_shot = params.include_one_shot.as_deref() == Some("true");
    let place = params.place()?;
    let db = state.db.clone();

    let stats = tokio::task::spawn_blocking(move || db.get_stats(include_one_shot, &place))
        .await
        .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
        .map_err(|e| internal_error(&e.to_string()))?;

    Ok(Json(serde_json::json!({
        "total_sessions": stats.session_count,
        "total_messages": stats.message_count
    })))
}

pub async fn activity(
    State(state): State<MachineWeb>,
    Query(params): Query<AnalyticsQuery>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    let granularity = params.granularity.as_deref().unwrap_or("day");
    let place = params.place()?;
    let db = state.db.clone();

    let series = tokio::task::spawn_blocking(move || {
        let conn = db.get_connection()?;
        let (sql, binds) = placed(
            "SELECT date(ended_at) as date, COUNT(*) as sessions, \
             SUM(message_count) as messages \
             FROM sessions WHERE deleted_at IS NULL AND ended_at IS NOT NULL{place} \
             GROUP BY date(ended_at) ORDER BY date ASC",
            &place,
        );
        let mut stmt = conn.prepare(&sql)?;
        let mut rows = stmt.query(bound(&binds).as_slice())?;

        let mut series = Vec::new();
        while let Some(row) = rows.next()? {
            series.push(serde_json::json!({
                "date": row.get::<&str, String>("date")?,
                "sessions": row.get::<&str, i64>("sessions")?,
                "messages": row.get::<&str, i64>("messages")?,
            }));
        }
        Ok(series)
    })
    .await
    .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
    .map_err(|e| internal_error(&e.to_string()))?;

    Ok(Json(serde_json::json!({
        "granularity": granularity,
        "series": series
    })))
}

pub async fn heatmap(
    State(state): State<MachineWeb>,
    Query(params): Query<AnalyticsQuery>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    let metric = params
        .metric
        .clone()
        .unwrap_or_else(|| "messages".to_string());
    let place = params.place()?;
    let db = state.db.clone();
    let metric_for_closure = metric.clone();

    let result = tokio::task::spawn_blocking(move || {
        let conn = db.get_connection()?;

        let value_col = if metric_for_closure == "sessions" {
            "COUNT(*)"
        } else {
            "SUM(message_count)"
        };

        let (sql, binds) = placed(
            &format!(
                "SELECT date(ended_at) as date, {} as value \
                 FROM sessions WHERE deleted_at IS NULL AND ended_at IS NOT NULL{{place}} \
                 GROUP BY date(ended_at) ORDER BY date ASC LIMIT 366",
                value_col
            ),
            &place,
        );

        let mut stmt = conn.prepare(&sql)?;
        let mut rows = stmt.query(bound(&binds).as_slice())?;

        let mut entries = Vec::new();
        let mut max_val: i64 = 0;

        while let Some(row) = rows.next()? {
            let val: i64 = row.get("value")?;
            if val > max_val {
                max_val = val;
            }
            entries.push(serde_json::json!({
                "date": row.get::<&str, String>("date")?,
                "value": val,
            }));
        }

        let (l1, l2, l3, l4) = if max_val > 0 {
            let q = max_val as f64 / 4.0;
            (
                (q * 0.25).ceil() as i64,
                (q * 1.0).ceil() as i64,
                (q * 2.0).ceil() as i64,
                (q * 3.0).ceil() as i64,
            )
        } else {
            (1, 5, 12, 30)
        };

        let entries_with_levels: Vec<serde_json::Value> = entries
            .into_iter()
            .map(|e| {
                let val = e.get("value").and_then(|v| v.as_i64()).unwrap_or(0);
                let level = if val == 0 {
                    0
                } else if val <= l1 {
                    1
                } else if val <= l2 {
                    2
                } else if val <= l3 {
                    3
                } else {
                    4
                };
                serde_json::json!({
                    "date": e.get("date"),
                    "value": val,
                    "level": level
                })
            })
            .collect();

        let entries_from = entries_with_levels
            .first()
            .and_then(|e| e.get("date"))
            .and_then(|d| d.as_str())
            .map(|s| s.to_string())
            .unwrap_or_default();

        Ok((entries_with_levels, entries_from, l1, l2, l3, l4))
    })
    .await
    .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
    .map_err(|e| internal_error(&e.to_string()))?;

    let (entries, entries_from, l1, l2, l3, l4) = result;

    Ok(Json(serde_json::json!({
        "metric": metric,
        "entries_from": entries_from,
        "levels": { "l1": l1, "l2": l2, "l3": l3, "l4": l4 },
        "entries": entries
    })))
}

/// Sessions by the project and worktree they ran in.
pub async fn projects(
    State(state): State<MachineWeb>,
    Query(params): Query<AnalyticsQuery>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    let place = params.place()?;
    let db = state.db.clone();

    let projects = tokio::task::spawn_blocking(move || {
        let conn = db.get_connection()?;
        let (sql, binds) = placed(
            "SELECT project, worktree, COUNT(*) as sessions, \
             SUM(message_count) as messages, \
             MIN(started_at) as first_session, MAX(ended_at) as last_session, \
             AVG(message_count) as avg_messages \
             FROM sessions WHERE deleted_at IS NULL{place} \
             GROUP BY project, worktree ORDER BY messages DESC",
            &place,
        );
        let mut stmt = conn.prepare(&sql)?;
        let mut rows = stmt.query(bound(&binds).as_slice())?;

        let mut projects = Vec::new();
        while let Some(row) = rows.next()? {
            projects.push(serde_json::json!({
                "project": row.get::<&str, Option<String>>("project")?,
                "worktree": row.get::<&str, Option<String>>("worktree")?,
                "sessions": row.get::<&str, i64>("sessions")?,
                "messages": row.get::<&str, i64>("messages")?,
                "first_session": row.get::<&str, Option<String>>("first_session")?,
                "last_session": row.get::<&str, Option<String>>("last_session")?,
                "avg_messages": row.get::<&str, f64>("avg_messages")?,
                "median_messages": 0,
                "agents": {},
                "daily_trend": 0.0
            }));
        }
        Ok(projects)
    })
    .await
    .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
    .map_err(|e| internal_error(&e.to_string()))?;

    Ok(Json(serde_json::json!({ "projects": projects })))
}

pub async fn hour_of_week(
    State(state): State<MachineWeb>,
    Query(params): Query<AnalyticsQuery>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    let place = params.place()?;
    let db = state.db.clone();

    let cells = tokio::task::spawn_blocking(move || {
        let conn = db.get_connection()?;
        let (sql, binds) = placed(
            "SELECT CAST(strftime('%w', ended_at) AS INTEGER) as dow, \
             CAST(strftime('%H', ended_at) AS INTEGER) as hour, \
             SUM(message_count) as messages \
             FROM sessions WHERE deleted_at IS NULL AND ended_at IS NOT NULL{place} \
             GROUP BY dow, hour",
            &place,
        );
        let mut stmt = conn.prepare(&sql)?;
        let mut rows = stmt.query(bound(&binds).as_slice())?;

        let mut cells = Vec::new();
        while let Some(row) = rows.next()? {
            let sqlite_dow: i64 = row.get("dow")?;
            let day_of_week = if sqlite_dow == 0 { 6 } else { sqlite_dow - 1 };
            cells.push(serde_json::json!({
                "day_of_week": day_of_week,
                "hour": row.get::<&str, i64>("hour")?,
                "messages": row.get::<&str, i64>("messages")?
            }));
        }
        Ok(cells)
    })
    .await
    .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
    .map_err(|e| internal_error(&e.to_string()))?;

    Ok(Json(serde_json::json!({ "cells": cells })))
}

pub async fn top_sessions(
    State(state): State<MachineWeb>,
    Query(params): Query<AnalyticsQuery>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    let metric = params
        .metric
        .clone()
        .unwrap_or_else(|| "messages".to_string());
    let place = params.place()?;
    let db = state.db.clone();
    let metric_for_query = metric.clone();

    let sessions = tokio::task::spawn_blocking(move || {
        let conn = db.get_connection()?;

        let order_col = if metric_for_query == "duration" {
            "(julianday(ended_at) - julianday(started_at)) * 24 * 60"
        } else {
            "message_count"
        };

        let (sql, binds) = placed(
            &format!(
                "SELECT id, project, worktree, first_message, message_count, \
                 CAST((julianday(ended_at) - julianday(started_at)) * 24 * 60 AS REAL) as duration_min \
                 FROM sessions WHERE deleted_at IS NULL AND started_at IS NOT NULL \
                 AND ended_at IS NOT NULL{{place}} \
                 ORDER BY {} DESC LIMIT 10",
                order_col
            ),
            &place,
        );

        let mut stmt = conn.prepare(&sql)?;
        let mut rows = stmt.query(bound(&binds).as_slice())?;

        let mut sessions = Vec::new();
        while let Some(row) = rows.next()? {
            sessions.push(serde_json::json!({
                "id": row.get::<&str, String>("id")?,
                "project": row.get::<&str, Option<String>>("project")?,
                "worktree": row.get::<&str, Option<String>>("worktree")?,
                "first_message": row.get::<&str, Option<String>>("first_message")?,
                "message_count": row.get::<&str, i64>("message_count")?,
                "duration_min": row.get::<&str, f64>("duration_min")?
            }));
        }
        Ok(sessions)
    }).await.unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
    .map_err(|e| internal_error(&e.to_string()))?;

    Ok(Json(serde_json::json!({
        "metric": metric,
        "sessions": sessions
    })))
}

#[derive(Deserialize)]
pub struct UsageQuery {
    pub from: Option<String>,
    pub to: Option<String>,
    pub project: Option<String>,
    pub worktree: Option<String>,
    pub placed: Option<String>,
    pub agent: Option<String>,
    pub model: Option<String>,
}

#[derive(serde::Serialize)]
pub struct UsageResponse {
    pub summary: UsageSummary,
    pub pricing_version: String,
    pub pricing_count: i64,
}

#[derive(serde::Serialize)]
pub struct DailyUsageResponse {
    pub daily: Vec<DailyUsage>,
    pub pricing_version: String,
}

#[derive(serde::Serialize)]
pub struct ModelUsageResponse {
    pub by_model: Vec<ModelUsage>,
    pub pricing_version: String,
}

#[derive(serde::Serialize)]
pub struct ProjectUsageResponse {
    pub by_project: Vec<ProjectUsage>,
    pub pricing_version: String,
}

pub async fn usage_summary(
    State(state): State<MachineWeb>,
    Query(params): Query<UsageQuery>,
) -> Result<Json<UsageResponse>, (StatusCode, Json<ErrorResponse>)> {
    use wsb::pricing::{fallback_pricing, find_pricing, FALLBACK_VERSION};

    let filter = UsageFilter {
        date_from: params.from,
        date_to: params.to,
        place: place_filter(params.project, params.worktree, params.placed)?,
        agent: params.agent,
        model: params.model,
    };
    let filter_clone = filter.clone();

    let db = state.db.clone();
    let db_summary = db.clone();

    let summary = tokio::task::spawn_blocking(move || {
        db_summary.get_usage_summary(&filter_clone)
    })
    .await
    .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
    .map_err(|e| internal_error(&e.to_string()))?;

    let db_pricing = tokio::task::spawn_blocking(move || db.get_model_pricing())
        .await
        .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
        .map_err(|e| internal_error(&e.to_string()))?;

    let pricing = if db_pricing.is_empty() {
        fallback_pricing()
    } else {
        db_pricing
    };

    let pricing_count = pricing.len() as i64;
    let estimated_cost = if let Some(ref model_filter) = filter.model {
        if let Some(p) = find_pricing(model_filter, &pricing) {
            let input_cost = (summary.total_input_tokens as f64 / 1_000_000.0) * p.input_per_mtok;
            let output_cost =
                (summary.total_output_tokens as f64 / 1_000_000.0) * p.output_per_mtok;
            input_cost + output_cost
        } else {
            0.0
        }
    } else {
        0.0
    };

    Ok(Json(UsageResponse {
        summary: UsageSummary {
            total_input_tokens: summary.total_input_tokens,
            total_output_tokens: summary.total_output_tokens,
            total_cache_creation_tokens: summary.total_cache_creation_tokens,
            total_cache_read_tokens: summary.total_cache_read_tokens,
            estimated_cost_usd: estimated_cost,
        },
        pricing_version: FALLBACK_VERSION.to_string(),
        pricing_count,
    }))
}

pub async fn usage_daily(
    State(state): State<MachineWeb>,
    Query(params): Query<UsageQuery>,
) -> Result<Json<DailyUsageResponse>, (StatusCode, Json<ErrorResponse>)> {
    use wsb::pricing::FALLBACK_VERSION;

    let filter = UsageFilter {
        date_from: params.from,
        date_to: params.to,
        place: place_filter(params.project, params.worktree, params.placed)?,
        agent: None,
        model: None,
    };

    let db = state.db.clone();

    let daily = tokio::task::spawn_blocking(move || db.get_daily_usage(&filter))
        .await
        .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
        .map_err(|e| internal_error(&e.to_string()))?;

    Ok(Json(DailyUsageResponse {
        daily,
        pricing_version: FALLBACK_VERSION.to_string(),
    }))
}

pub async fn usage_by_model(
    State(state): State<MachineWeb>,
    Query(params): Query<UsageQuery>,
) -> Result<Json<ModelUsageResponse>, (StatusCode, Json<ErrorResponse>)> {
    use wsb::pricing::{fallback_pricing, find_pricing, FALLBACK_VERSION};

    let filter = UsageFilter {
        date_from: params.from,
        date_to: params.to,
        place: place_filter(params.project, params.worktree, params.placed)?,
        agent: None,
        model: None,
    };

    let db = state.db.clone();
    let db_by_model = db.clone();

    let by_model =
        tokio::task::spawn_blocking(move || db_by_model.get_usage_by_model(&filter))
            .await
            .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
            .map_err(|e| internal_error(&e.to_string()))?;

    let db_pricing = tokio::task::spawn_blocking(move || db.get_model_pricing())
        .await
        .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
        .unwrap_or_default();

    let pricing = if db_pricing.is_empty() {
        fallback_pricing()
    } else {
        db_pricing
    };

    let by_model_with_cost: Vec<ModelUsage> = by_model
        .into_iter()
        .map(|mut m| {
            if let Some(p) = find_pricing(&m.model, &pricing) {
                let input_cost = (m.input_tokens as f64 / 1_000_000.0) * p.input_per_mtok;
                let output_cost = (m.output_tokens as f64 / 1_000_000.0) * p.output_per_mtok;
                m.estimated_cost_usd = input_cost + output_cost;
            }
            m
        })
        .collect();

    Ok(Json(ModelUsageResponse {
        by_model: by_model_with_cost,
        pricing_version: FALLBACK_VERSION.to_string(),
    }))
}

pub async fn usage_by_project(
    State(state): State<MachineWeb>,
    Query(params): Query<UsageQuery>,
) -> Result<Json<ProjectUsageResponse>, (StatusCode, Json<ErrorResponse>)> {
    use wsb::pricing::FALLBACK_VERSION;

    let filter = UsageFilter {
        date_from: params.from,
        date_to: params.to,
        place: place_filter(params.project, params.worktree, params.placed)?,
        agent: None,
        model: None,
    };

    let db = state.db.clone();

    let by_project =
        tokio::task::spawn_blocking(move || db.get_usage_by_project(&filter))
            .await
            .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
            .map_err(|e| internal_error(&e.to_string()))?;

    Ok(Json(ProjectUsageResponse {
        by_project,
        pricing_version: FALLBACK_VERSION.to_string(),
    }))
}

pub async fn sync_pricing(
    State(state): State<MachineWeb>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    let db = state.db.clone();
    let db_count = db.clone();

    tokio::task::spawn_blocking(move || db.sync_fallback_pricing())
        .await
        .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
        .map_err(|e| internal_error(&e.to_string()))?;

    let count = tokio::task::spawn_blocking(move || db_count.get_model_pricing())
        .await
        .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
        .map_err(|e| internal_error(&e.to_string()))?
        .len();

    Ok(Json(serde_json::json!({
        "synced": count,
        "version": wsb::pricing::FALLBACK_VERSION
    })))
}