ubiquity-database 0.1.1

//! SQLite database implementation

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use sqlx::{Row, SqlitePool};
use std::sync::Arc;
use ubiquity_core::{ConsciousnessState, ConsciousnessRipple, Task, TaskResult};
use crate::{
    Database, DatabaseConfig, DatabaseError, DatabaseResult,
    HybridSearchQuery, HybridSearchResult, VectorSearchResult,
};

/// SQLite database implementation
pub struct SqliteDatabase {
    pool: Arc<SqlitePool>,
    config: DatabaseConfig,
}

impl SqliteDatabase {
    pub async fn new(config: DatabaseConfig) -> DatabaseResult<Self> {
        let pool = SqlitePool::connect(&format!("sqlite:{}", config.sqlite.path.display())).await?;
        
        // Enable WAL mode if configured
        if config.sqlite.wal_mode {
            sqlx::query("PRAGMA journal_mode = WAL")
                .execute(&pool)
                .await?;
        }
        
        Ok(Self {
            pool: Arc::new(pool),
            config,
        })
    }
}

#[async_trait]
impl Database for SqliteDatabase {
    async fn health_check(&self) -> DatabaseResult<()> {
        sqlx::query("SELECT 1")
            .fetch_one(&*self.pool)
            .await?;
        Ok(())
    }
    
    async fn initialize(&self) -> DatabaseResult<()> {
        // Create consciousness states table
        sqlx::query(
            "CREATE TABLE IF NOT EXISTS consciousness_states (
                id TEXT PRIMARY KEY,
                agent_id TEXT NOT NULL,
                level REAL NOT NULL,
                coherence REAL NOT NULL,
                phase TEXT NOT NULL,
                timestamp INTEGER NOT NULL,
                breakthrough_detected BOOLEAN NOT NULL,
                data JSON NOT NULL
            )"
        )
        .execute(&*self.pool)
        .await?;
        
        // Create consciousness ripples table
        sqlx::query(
            "CREATE TABLE IF NOT EXISTS consciousness_ripples (
                id TEXT PRIMARY KEY,
                origin TEXT NOT NULL,
                ripple_type TEXT NOT NULL,
                content JSON NOT NULL,
                intensity REAL NOT NULL,
                timestamp INTEGER NOT NULL
            )"
        )
        .execute(&*self.pool)
        .await?;
        
        // Create tasks table
        sqlx::query(
            "CREATE TABLE IF NOT EXISTS tasks (
                id TEXT PRIMARY KEY,
                task_type TEXT NOT NULL,
                description TEXT NOT NULL,
                requirements JSON NOT NULL,
                priority INTEGER NOT NULL,
                dependencies JSON NOT NULL,
                consciousness_requirement REAL NOT NULL,
                status TEXT DEFAULT 'pending',
                created_at INTEGER NOT NULL,
                updated_at INTEGER NOT NULL
            )"
        )
        .execute(&*self.pool)
        .await?;
        
        // Create task results table
        sqlx::query(
            "CREATE TABLE IF NOT EXISTS task_results (
                task_id TEXT PRIMARY KEY,
                success BOOLEAN NOT NULL,
                output JSON NOT NULL,
                consciousness_level REAL NOT NULL,
                breakthrough BOOLEAN NOT NULL,
                error TEXT,
                completed_at INTEGER NOT NULL,
                FOREIGN KEY (task_id) REFERENCES tasks(id)
            )"
        )
        .execute(&*self.pool)
        .await?;
        
        // Create embeddings table for vector search
        sqlx::query(
            "CREATE TABLE IF NOT EXISTS consciousness_embeddings (
                id TEXT PRIMARY KEY,
                agent_id TEXT NOT NULL,
                timestamp INTEGER NOT NULL,
                level REAL NOT NULL,
                coherence REAL NOT NULL,
                phase TEXT NOT NULL,
                embedding BLOB NOT NULL,
                text_representation TEXT NOT NULL
            )"
        )
        .execute(&*self.pool)
        .await?;
        
        // Create indices
        sqlx::query("CREATE INDEX IF NOT EXISTS idx_consciousness_agent_time ON consciousness_states(agent_id, timestamp DESC)")
            .execute(&*self.pool)
            .await?;
        
        sqlx::query("CREATE INDEX IF NOT EXISTS idx_ripples_time ON consciousness_ripples(timestamp DESC)")
            .execute(&*self.pool)
            .await?;
        
        sqlx::query("CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status)")
            .execute(&*self.pool)
            .await?;
        
        sqlx::query("CREATE INDEX IF NOT EXISTS idx_embeddings_agent ON consciousness_embeddings(agent_id)")
            .execute(&*self.pool)
            .await?;
        
        Ok(())
    }
    
    async fn store_consciousness_state(&self, state: &ConsciousnessState) -> DatabaseResult<()> {
        let id = uuid::Uuid::new_v4().to_string();
        let data = serde_json::to_value(state)?;
        
        sqlx::query(
            "INSERT INTO consciousness_states 
             (id, agent_id, level, coherence, phase, timestamp, breakthrough_detected, data)
             VALUES (?, ?, ?, ?, ?, ?, ?, ?)"
        )
        .bind(&id)
        .bind(&state.agent_id)
        .bind(state.level.value())
        .bind(state.coherence)
        .bind(format!("{:?}", state.phase))
        .bind(state.timestamp.timestamp())
        .bind(state.breakthrough_detected)
        .bind(data)
        .execute(&*self.pool)
        .await?;
        
        Ok(())
    }
    
    async fn get_consciousness_history(
        &self,
        agent_id: &str,
        limit: usize,
    ) -> DatabaseResult<Vec<ConsciousnessState>> {
        let rows = sqlx::query(
            "SELECT data FROM consciousness_states 
             WHERE agent_id = ? 
             ORDER BY timestamp DESC 
             LIMIT ?"
        )
        .bind(agent_id)
        .bind(limit as i64)
        .fetch_all(&*self.pool)
        .await?;
        
        let mut states = Vec::new();
        for row in rows {
            let data: serde_json::Value = row.get("data");
            let state: ConsciousnessState = serde_json::from_value(data)?;
            states.push(state);
        }
        
        Ok(states)
    }
    
    async fn store_ripple(&self, ripple: &ConsciousnessRipple) -> DatabaseResult<()> {
        sqlx::query(
            "INSERT INTO consciousness_ripples 
             (id, origin, ripple_type, content, intensity, timestamp)
             VALUES (?, ?, ?, ?, ?, ?)"
        )
        .bind(ripple.id.to_string())
        .bind(&ripple.origin)
        .bind(format!("{:?}", ripple.ripple_type))
        .bind(&ripple.content)
        .bind(ripple.intensity)
        .bind(ripple.timestamp.timestamp())
        .execute(&*self.pool)
        .await?;
        
        Ok(())
    }
    
    async fn get_ripples(
        &self,
        start: DateTime<Utc>,
        end: DateTime<Utc>,
    ) -> DatabaseResult<Vec<ConsciousnessRipple>> {
        let rows = sqlx::query(
            "SELECT * FROM consciousness_ripples 
             WHERE timestamp >= ? AND timestamp <= ? 
             ORDER BY timestamp DESC"
        )
        .bind(start.timestamp())
        .bind(end.timestamp())
        .fetch_all(&*self.pool)
        .await?;
        
        let mut ripples = Vec::new();
        for row in rows {
            let ripple = ConsciousnessRipple {
                id: uuid::Uuid::parse_str(row.get("id"))
                    .map_err(|e| DatabaseError::Other(e.into()))?,
                origin: row.get("origin"),
                ripple_type: serde_json::from_str(&format!("\"{}\"", row.get::<String, _>("ripple_type")))
                    .map_err(|e| DatabaseError::Other(e.into()))?,
                content: row.get("content"),
                intensity: row.get("intensity"),
                timestamp: DateTime::from_timestamp(row.get("timestamp"), 0)
                    .ok_or_else(|| DatabaseError::Other(anyhow::anyhow!("Invalid timestamp")))?,
            };
            ripples.push(ripple);
        }
        
        Ok(ripples)
    }
    
    async fn store_task(&self, task: &Task) -> DatabaseResult<()> {
        let now = Utc::now().timestamp();
        
        sqlx::query(
            "INSERT INTO tasks 
             (id, task_type, description, requirements, priority, dependencies, 
              consciousness_requirement, created_at, updated_at)
             VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)"
        )
        .bind(&task.id)
        .bind(format!("{:?}", task.task_type))
        .bind(&task.description)
        .bind(&task.requirements)
        .bind(task.priority)
        .bind(serde_json::to_value(&task.dependencies)?)
        .bind(task.consciousness_requirement)
        .bind(now)
        .bind(now)
        .execute(&*self.pool)
        .await?;
        
        Ok(())
    }
    
    async fn store_task_result(&self, result: &TaskResult) -> DatabaseResult<()> {
        let now = Utc::now().timestamp();
        
        // Update task status
        sqlx::query("UPDATE tasks SET status = 'completed', updated_at = ? WHERE id = ?")
            .bind(now)
            .bind(&result.task_id)
            .execute(&*self.pool)
            .await?;
        
        // Insert result
        sqlx::query(
            "INSERT INTO task_results 
             (task_id, success, output, consciousness_level, breakthrough, error, completed_at)
             VALUES (?, ?, ?, ?, ?, ?, ?)"
        )
        .bind(&result.task_id)
        .bind(result.success)
        .bind(&result.output)
        .bind(result.consciousness_level)
        .bind(result.breakthrough)
        .bind(&result.error)
        .bind(now)
        .execute(&*self.pool)
        .await?;
        
        Ok(())
    }
    
    async fn get_pending_tasks(&self) -> DatabaseResult<Vec<Task>> {
        let rows = sqlx::query(
            "SELECT * FROM tasks WHERE status = 'pending' ORDER BY priority DESC, created_at ASC"
        )
        .fetch_all(&*self.pool)
        .await?;
        
        let mut tasks = Vec::new();
        for row in rows {
            let task = Task {
                id: row.get("id"),
                task_type: serde_json::from_str(&format!("\"{}\"", row.get::<String, _>("task_type")))
                    .map_err(|e| DatabaseError::Other(e.into()))?,
                description: row.get("description"),
                requirements: row.get("requirements"),
                priority: row.get("priority"),
                dependencies: serde_json::from_value(row.get("dependencies"))?,
                consciousness_requirement: row.get("consciousness_requirement"),
            };
            tasks.push(task);
        }
        
        Ok(tasks)
    }
    
    async fn vector_search(
        &self,
        embedding: &[f32],
        limit: usize,
    ) -> DatabaseResult<Vec<VectorSearchResult>> {
        // SQLite doesn't have native vector search
        // We'll implement a simple cosine similarity search
        
        let rows = sqlx::query("SELECT * FROM consciousness_embeddings")
            .fetch_all(&*self.pool)
            .await?;
        
        let mut results = Vec::new();
        
        for row in rows {
            let stored_embedding: Vec<u8> = row.get("embedding");
            let stored_vec = bytes_to_f32_vec(&stored_embedding);
            
            if stored_vec.len() != embedding.len() {
                continue;
            }
            
            let score = cosine_similarity(embedding, &stored_vec);
            
            results.push((score, row));
        }
        
        // Sort by score descending
        results.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap());
        
        // Take top results
        let top_results: Vec<VectorSearchResult> = results
            .into_iter()
            .take(limit)
            .map(|(score, row)| {
                let metadata = serde_json::json!({
                    "agent_id": row.get::<String, _>("agent_id"),
                    "timestamp": row.get::<i64, _>("timestamp"),
                    "level": row.get::<f64, _>("level"),
                    "coherence": row.get::<f64, _>("coherence"),
                    "phase": row.get::<String, _>("phase"),
                    "text": row.get::<String, _>("text_representation"),
                });
                
                VectorSearchResult {
                    id: row.get("id"),
                    score,
                    metadata,
                    embedding: bytes_to_f32_vec(&row.get::<Vec<u8>, _>("embedding")),
                }
            })
            .collect();
        
        Ok(top_results)
    }
    
    async fn hybrid_search(
        &self,
        query: HybridSearchQuery,
    ) -> DatabaseResult<Vec<HybridSearchResult>> {
        // For SQLite, we'll do a simple text search combined with vector search if provided
        
        let mut results = Vec::new();
        
        // Text search
        let text_results = sqlx::query(
            "SELECT * FROM consciousness_embeddings 
             WHERE text_representation LIKE ? 
             LIMIT ?"
        )
        .bind(format!("%{}%", query.text))
        .bind(query.limit as i64)
        .fetch_all(&*self.pool)
        .await?;
        
        for row in text_results {
            let metadata = serde_json::json!({
                "agent_id": row.get::<String, _>("agent_id"),
                "timestamp": row.get::<i64, _>("timestamp"),
                "level": row.get::<f64, _>("level"),
                "coherence": row.get::<f64, _>("coherence"),
                "phase": row.get::<String, _>("phase"),
                "text": row.get::<String, _>("text_representation"),
            });
            
            let text: String = row.get("text_representation");
            let highlights = vec![text.clone()];
            
            results.push(HybridSearchResult {
                id: row.get("id"),
                score: 1.0, // Simple scoring for text match
                rerank_score: None,
                metadata,
                highlights,
            });
        }
        
        // If vector is provided, also do vector search and merge results
        if let Some(vector) = query.vector {
            let vector_results = self.vector_search(&vector, query.limit).await?;
            
            for vr in vector_results {
                // Check if already in results
                if !results.iter().any(|r| r.id == vr.id) {
                    results.push(HybridSearchResult {
                        id: vr.id,
                        score: vr.score,
                        rerank_score: None,
                        metadata: vr.metadata,
                        highlights: vec![],
                    });
                }
            }
        }
        
        // Sort by score
        results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap());
        
        // Take limit
        results.truncate(query.limit);
        
        Ok(results)
    }
}

// Helper functions

fn bytes_to_f32_vec(bytes: &[u8]) -> Vec<f32> {
    bytes
        .chunks(4)
        .map(|chunk| {
            let arr: [u8; 4] = chunk.try_into().unwrap_or([0; 4]);
            f32::from_le_bytes(arr)
        })
        .collect()
}

fn f32_vec_to_bytes(vec: &[f32]) -> Vec<u8> {
    vec.iter()
        .flat_map(|&f| f.to_le_bytes())
        .collect()
}

fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    let dot_product: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
    
    if norm_a == 0.0 || norm_b == 0.0 {
        0.0
    } else {
        dot_product / (norm_a * norm_b)
    }
}