systemprompt-analytics 0.2.2

use anyhow::Result;
use chrono::{DateTime, Utc};
use sqlx::PgPool;
use std::sync::Arc;
use systemprompt_database::DbPool;

use crate::models::cli::{CostBreakdownRow, CostSummaryRow, CostTrendRow, PreviousCostRow};

#[derive(Debug)]
pub struct CostAnalyticsRepository {
    pool: Arc<PgPool>,
}

impl CostAnalyticsRepository {
    pub fn new(db: &DbPool) -> Result<Self> {
        let pool = db.pool_arc()?;
        Ok(Self { pool })
    }

    pub async fn get_summary(
        &self,
        start: DateTime<Utc>,
        end: DateTime<Utc>,
    ) -> Result<CostSummaryRow> {
        sqlx::query_as!(
            CostSummaryRow,
            r#"
            SELECT
                COUNT(*)::bigint as "total_requests!",
                SUM(cost_microdollars)::bigint as "total_cost",
                SUM(tokens_used)::bigint as "total_tokens"
            FROM ai_requests
            WHERE created_at >= $1 AND created_at < $2
            "#,
            start,
            end
        )
        .fetch_one(&*self.pool)
        .await
        .map_err(Into::into)
    }

    pub async fn get_previous_cost(
        &self,
        start: DateTime<Utc>,
        end: DateTime<Utc>,
    ) -> Result<PreviousCostRow> {
        sqlx::query_as!(
            PreviousCostRow,
            r#"
            SELECT SUM(cost_microdollars)::bigint as "cost"
            FROM ai_requests
            WHERE created_at >= $1 AND created_at < $2
            "#,
            start,
            end
        )
        .fetch_one(&*self.pool)
        .await
        .map_err(Into::into)
    }

    pub async fn get_breakdown_by_model(
        &self,
        start: DateTime<Utc>,
        end: DateTime<Utc>,
        limit: i64,
    ) -> Result<Vec<CostBreakdownRow>> {
        sqlx::query_as!(
            CostBreakdownRow,
            r#"
            SELECT
                model as "name!",
                COALESCE(SUM(cost_microdollars), 0)::bigint as "cost!",
                COUNT(*)::bigint as "requests!",
                COALESCE(SUM(tokens_used), 0)::bigint as "tokens!"
            FROM ai_requests
            WHERE created_at >= $1 AND created_at < $2
            GROUP BY model
            ORDER BY SUM(cost_microdollars) DESC NULLS LAST
            LIMIT $3
            "#,
            start,
            end,
            limit
        )
        .fetch_all(&*self.pool)
        .await
        .map_err(Into::into)
    }

    pub async fn get_breakdown_by_provider(
        &self,
        start: DateTime<Utc>,
        end: DateTime<Utc>,
        limit: i64,
    ) -> Result<Vec<CostBreakdownRow>> {
        sqlx::query_as!(
            CostBreakdownRow,
            r#"
            SELECT
                provider as "name!",
                COALESCE(SUM(cost_microdollars), 0)::bigint as "cost!",
                COUNT(*)::bigint as "requests!",
                COALESCE(SUM(tokens_used), 0)::bigint as "tokens!"
            FROM ai_requests
            WHERE created_at >= $1 AND created_at < $2
            GROUP BY provider
            ORDER BY SUM(cost_microdollars) DESC NULLS LAST
            LIMIT $3
            "#,
            start,
            end,
            limit
        )
        .fetch_all(&*self.pool)
        .await
        .map_err(Into::into)
    }

    pub async fn get_breakdown_by_agent(
        &self,
        start: DateTime<Utc>,
        end: DateTime<Utc>,
        limit: i64,
    ) -> Result<Vec<CostBreakdownRow>> {
        sqlx::query_as!(
            CostBreakdownRow,
            r#"
            (
                SELECT
                    at.agent_name as "name!",
                    COALESCE(SUM(r.cost_microdollars), 0)::bigint as "cost!",
                    COUNT(*)::bigint as "requests!",
                    COALESCE(SUM(r.tokens_used), 0)::bigint as "tokens!"
                FROM ai_requests r
                INNER JOIN agent_tasks at ON at.task_id = r.task_id
                WHERE r.created_at >= $1 AND r.created_at < $2
                  AND at.agent_name IS NOT NULL
                GROUP BY at.agent_name
                ORDER BY SUM(r.cost_microdollars) DESC NULLS LAST
                LIMIT $3
            )
            UNION ALL
            (
                SELECT
                    'unattributed' as "name!",
                    COALESCE(SUM(r.cost_microdollars), 0)::bigint as "cost!",
                    COUNT(*)::bigint as "requests!",
                    COALESCE(SUM(r.tokens_used), 0)::bigint as "tokens!"
                FROM ai_requests r
                LEFT JOIN agent_tasks at ON at.task_id = r.task_id
                WHERE r.created_at >= $1 AND r.created_at < $2
                  AND (r.task_id IS NULL OR at.agent_name IS NULL)
                HAVING COUNT(*) > 0
            )
            "#,
            start,
            end,
            limit
        )
        .fetch_all(&*self.pool)
        .await
        .map_err(Into::into)
    }

    pub async fn get_costs_for_trends(
        &self,
        start: DateTime<Utc>,
        end: DateTime<Utc>,
    ) -> Result<Vec<CostTrendRow>> {
        sqlx::query_as!(
            CostTrendRow,
            r#"
            SELECT
                created_at as "created_at!",
                cost_microdollars,
                tokens_used
            FROM ai_requests
            WHERE created_at >= $1 AND created_at < $2
            ORDER BY created_at
            "#,
            start,
            end
        )
        .fetch_all(&*self.pool)
        .await
        .map_err(Into::into)
    }
}