tusk/src-tauri/src/models/ai.rs

use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "lowercase")]
pub enum AiProvider {
    #[default]
    Ollama,
    OpenAi,
    Anthropic,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AiSettings {
    pub provider: AiProvider,
    pub ollama_url: String,
    pub openai_api_key: Option<String>,
    pub anthropic_api_key: Option<String>,
    pub model: String,
}

impl Default for AiSettings {
    fn default() -> Self {
        Self {
            provider: AiProvider::Ollama,
            ollama_url: "http://localhost:11434".to_string(),
            openai_api_key: None,
            anthropic_api_key: None,
            model: String::new(),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OllamaChatMessage {
    pub role: String,
    pub content: String,
}

#[derive(Debug, Clone, Serialize)]
pub struct OllamaChatRequest {
    pub model: String,
    pub messages: Vec<OllamaChatMessage>,
    pub stream: bool,
}

#[derive(Debug, Deserialize)]
pub struct OllamaChatResponse {
    pub message: OllamaChatMessage,
}

#[derive(Debug, Deserialize)]
pub struct OllamaTagsResponse {
    pub models: Vec<OllamaModel>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OllamaModel {
    pub name: String,
}

// --- Wave 1: Validation ---

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ValidationStatus {
    Pending,
    Generating,
    Running,
    Passed,
    Failed,
    Error,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationRule {
    pub id: String,
    pub description: String,
    pub generated_sql: String,
    pub status: ValidationStatus,
    pub violation_count: u64,
    pub sample_violations: Vec<Vec<serde_json::Value>>,
    pub violation_columns: Vec<String>,
    pub error: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationReport {
    pub rules: Vec<ValidationRule>,
    pub total_rules: usize,
    pub passed: usize,
    pub failed: usize,
    pub errors: usize,
    pub execution_time_ms: u128,
}

// --- Wave 2: Data Generator ---

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GenerateDataParams {
    pub connection_id: String,
    pub schema: String,
    pub table: String,
    pub row_count: u32,
    pub include_related: bool,
    pub custom_instructions: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GeneratedDataPreview {
    pub tables: Vec<GeneratedTableData>,
    pub insert_order: Vec<String>,
    pub total_rows: u32,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GeneratedTableData {
    pub schema: String,
    pub table: String,
    pub columns: Vec<String>,
    pub rows: Vec<Vec<serde_json::Value>>,
    pub row_count: u32,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DataGenProgress {
    pub gen_id: String,
    pub stage: String,
    pub percent: u8,
    pub message: String,
    pub detail: Option<String>,
}

// --- Wave 3A: Index Advisor ---

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TableStats {
    pub schema: String,
    pub table: String,
    pub seq_scan: i64,
    pub idx_scan: i64,
    pub n_live_tup: i64,
    pub table_size: String,
    pub index_size: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexStats {
    pub schema: String,
    pub table: String,
    pub index_name: String,
    pub idx_scan: i64,
    pub index_size: String,
    pub definition: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SlowQuery {
    pub query: String,
    pub calls: i64,
    pub total_time_ms: f64,
    pub mean_time_ms: f64,
    pub rows: i64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum IndexRecommendationType {
    CreateIndex,
    DropIndex,
    ReplaceIndex,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexRecommendation {
    pub id: String,
    pub recommendation_type: IndexRecommendationType,
    pub table_schema: String,
    pub table_name: String,
    pub index_name: Option<String>,
    pub ddl: String,
    pub rationale: String,
    pub estimated_impact: String,
    pub priority: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexAdvisorReport {
    pub table_stats: Vec<TableStats>,
    pub index_stats: Vec<IndexStats>,
    pub slow_queries: Vec<SlowQuery>,
    pub recommendations: Vec<IndexRecommendation>,
    pub has_pg_stat_statements: bool,
}