iAs/src/db/models.rs

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sqlx::PgPool;

use crate::state::AuthState;

// ─── 认证状态 ───

/// 从数据库加载认证状态
pub async fn load_auth(pool: &PgPool) -> Option<AuthState> {
    let row = sqlx::query_as::<_, (serde_json::Value,)>(
        "SELECT value FROM app_state WHERE key = 'auth'",
    )
    .fetch_optional(pool)
    .await
    .ok()??;

    serde_json::from_value(row.0).ok()
}

/// 保存认证状态到数据库
pub async fn save_auth(pool: &PgPool, auth: &AuthState) -> Result<(), String> {
    let value = serde_json::to_value(auth).map_err(|e| format!("序列化 auth 失败: {}", e))?;

    sqlx::query(
        r#"
        INSERT INTO app_state (key, value, updated_at)
        VALUES ('auth', $1, NOW())
        ON CONFLICT (key) DO UPDATE
        SET value = EXCLUDED.value, updated_at = NOW()
        "#,
    )
    .bind(&value)
    .execute(pool)
    .await
    .map_err(|e| format!("保存 auth 到数据库失败: {}", e))?;

    Ok(())
}

// ─── 聊天记录 ───

#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
pub struct ChatRecord {
    pub id: i64,
    pub created_at: DateTime<Utc>,
    pub direction: String,
    pub user_id: String,
    pub account_id: String,
    pub text: String,
    pub source: String,
    pub context_token: String,
    pub message_id: String,
}

/// 插入一条聊天记录
pub async fn insert_chat_record(
    pool: &PgPool,
    direction: &str,
    user_id: &str,
    account_id: &str,
    text: &str,
    source: &str,
    context_token: Option<&str>,
    message_id: &str,
) -> Result<i64, String> {
    let ctx = context_token.unwrap_or("");

    let row: Option<(i64,)> = sqlx::query_as(
        r#"
        INSERT INTO chat_records (direction, user_id, account_id, text, source, context_token, message_id)
        VALUES ($1, $2, $3, $4, $5, $6, $7)
        ON CONFLICT (message_id) WHERE message_id NOT IN ('', '0') DO NOTHING
        RETURNING id
        "#,
    )
    .bind(direction)
    .bind(user_id)
    .bind(account_id)
    .bind(text)
    .bind(source)
    .bind(ctx)
    .bind(message_id)
    .fetch_optional(pool)
    .await
    .map_err(|e| format!("插入聊天记录失败: {}", e))?;

    Ok(row.map(|r| r.0).unwrap_or(0))
}

/// 查询最近的聊天记录（用户维度）
pub async fn list_recent_chat_records(
    pool: &PgPool,
    user_id: &str,
    limit: i64,
) -> Result<Vec<ChatRecord>, String> {
    let records = sqlx::query_as::<_, ChatRecord>(
        r#"
        SELECT id, created_at, direction, user_id, account_id, text, source, context_token, message_id
        FROM chat_records
        WHERE user_id = $1
        ORDER BY created_at DESC
        LIMIT $2
        "#,
    )
    .bind(user_id)
    .bind(limit)
    .fetch_all(pool)
    .await
    .map_err(|e| format!("查询聊天记录失败: {}", e))?;

    Ok(records)
}

// ─── LLM 用量 ───

/// 插入一条 LLM 用量记录
/// kind/entry 为 iPet 遗留兼容字段，写入默认值 ("llm_call" / null)
pub async fn insert_llm_usage(
    pool: &PgPool,
    user_id: &str,
    model: &str,
    provider: &str,
    prompt_tokens: u32,
    completion_tokens: u32,
    cache_hit_tokens: u32,
    cache_miss_tokens: u32,
) -> Result<(), String> {
    let total = prompt_tokens + completion_tokens;
    sqlx::query(
        r#"
        INSERT INTO llm_usage (user_id, model, provider, prompt_tokens, completion_tokens, total_tokens, cache_hit_tokens, cache_miss_tokens, kind, entry)
        VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
        "#,
    )
    .bind(user_id)
    .bind(model)
    .bind(provider)
    .bind(prompt_tokens as i32)
    .bind(completion_tokens as i32)
    .bind(total as i32)
    .bind(cache_hit_tokens as i32)
    .bind(cache_miss_tokens as i32)
    .bind("llm_call")
    .bind(serde_json::Value::Null)
    .execute(pool)
    .await
    .map_err(|e| format!("插入 LLM 用量失败: {}", e))?;
    Ok(())
}

/// LLM 用量聚合统计
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct LlmUsageStats {
    pub total_calls: i64,
    pub total_prompt_tokens: i64,
    pub total_completion_tokens: i64,
    pub total_tokens: i64,
    pub total_cache_hit: i64,
    pub total_cache_miss: i64,
}

/// 查询 LLM 用量统计（按时间范围过滤）
pub async fn query_llm_usage_stats(
    pool: &PgPool,
    since: Option<chrono::DateTime<chrono::Utc>>,
    until: Option<chrono::DateTime<chrono::Utc>>,
    model_filter: Option<&str>,
) -> Result<LlmUsageStats, String> {
    let since = since.unwrap_or_else(|| {
        chrono::Utc::now() - chrono::Duration::days(7)
    });
    let until = until.unwrap_or_else(chrono::Utc::now);

    let row = sqlx::query_as::<_, LlmUsageStats>(
        r#"
        SELECT
            COUNT(*)::bigint as total_calls,
            COALESCE(SUM(prompt_tokens), 0)::bigint as total_prompt_tokens,
            COALESCE(SUM(completion_tokens), 0)::bigint as total_completion_tokens,
            COALESCE(SUM(total_tokens), 0)::bigint as total_tokens,
            COALESCE(SUM(cache_hit_tokens), 0)::bigint as total_cache_hit,
            COALESCE(SUM(cache_miss_tokens), 0)::bigint as total_cache_miss
        FROM llm_usage
        WHERE created_at >= $1 AND created_at <= $2
          AND ($3::text IS NULL OR model = $3)
        "#,
    )
    .bind(since)
    .bind(until)
    .bind(model_filter)
    .fetch_one(pool)
    .await
    .map_err(|e| format!("查询 LLM 用量失败: {}", e))?;

    Ok(row)
}

// ─── 会话摘要 ───

/// 保存一条摘要到数据库
pub async fn save_summary(
    pool: &PgPool,
    user_id: &str,
    text: &str,
    reason: &str,
    message_count: i32,
) -> Result<(), String> {
    sqlx::query(
        "INSERT INTO session_summaries (user_id, summary_text, reason, message_count) VALUES ($1, $2, $3, $4)",
    )
    .bind(user_id)
    .bind(text)
    .bind(reason)
    .bind(message_count)
    .execute(pool)
    .await
    .map_err(|e| format!("保存摘要失败: {}", e))?;
    Ok(())
}

/// 加载最近的摘要（用于 read_summaries 工具），按 user_id 过滤
pub async fn load_summaries(pool: &PgPool, user_id: &str, limit: i64) -> Result<Vec<String>, String> {
    let rows = sqlx::query_as::<_, (String,)>(
        "SELECT summary_text FROM session_summaries WHERE user_id = $1 ORDER BY created_at DESC LIMIT $2",
    )
    .bind(user_id)
    .bind(limit)
    .fetch_all(pool)
    .await
    .map_err(|e| format!("查询摘要失败: {}", e))?;
    Ok(rows.into_iter().map(|(t,)| t).collect())
}