ubiquity-database 0.1.1

//! Astra DB implementation with full vector and hybrid search support

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::sync::Arc;
use ubiquity_core::{ConsciousnessState, ConsciousnessRipple, Task, TaskResult};
use crate::{
    Database, DatabaseConfig, DatabaseError, DatabaseResult,
    HybridSearchQuery, HybridSearchResult, VectorSearchResult,
    embeddings::{ConsciousnessEmbedding, create_embedding_generator},
};

/// Astra DB database implementation
pub struct AstraDatabase {
    client: Arc<reqwest::Client>,
    config: DatabaseConfig,
    embedding_generator: Box<dyn crate::embeddings::EmbeddingGenerator>,
}

impl AstraDatabase {
    pub async fn new(config: DatabaseConfig) -> DatabaseResult<Self> {
        let client = Arc::new(reqwest::Client::new());
        let embedding_generator = create_embedding_generator(&config.embeddings);
        
        Ok(Self {
            client,
            config,
            embedding_generator,
        })
    }
    
    /// Get base URL for Astra DB API
    fn api_url(&self) -> String {
        format!("{}/api/json/v1/{}", 
            self.config.astra.endpoint,
            self.config.astra.keyspace
        )
    }
    
    /// Get auth headers
    fn auth_headers(&self) -> reqwest::header::HeaderMap {
        let mut headers = reqwest::header::HeaderMap::new();
        headers.insert(
            "X-Cassandra-Token",
            self.config.astra.token.parse().unwrap(),
        );
        headers.insert(
            reqwest::header::CONTENT_TYPE,
            "application/json".parse().unwrap(),
        );
        headers
    }
}

#[async_trait]
impl Database for AstraDatabase {
    async fn health_check(&self) -> DatabaseResult<()> {
        let response = self.client
            .get(&format!("{}/health", self.api_url()))
            .headers(self.auth_headers())
            .send()
            .await
            .map_err(|e| DatabaseError::Connection(e.to_string()))?;
        
        if !response.status().is_success() {
            return Err(DatabaseError::Connection(
                format!("Health check failed: {}", response.status())
            ));
        }
        
        Ok(())
    }
    
    async fn initialize(&self) -> DatabaseResult<()> {
        let collections = &self.config.astra.collections;
        
        // Create consciousness states collection with vectorize
        self.create_collection(
            &collections.consciousness,
            json!({
                "defaultId": { "type": "uuid" },
                "vector": {
                    "dimension": self.config.embeddings.dimension,
                    "service": {
                        "provider": self.config.embeddings.provider,
                        "modelName": self.config.embeddings.model,
                    }
                },
                "indexing": {
                    "allow": ["agent_id", "level", "coherence", "phase", "timestamp"],
                    "default": true
                },
                "lexical": {
                    "analyzer": "standard"
                }
            })
        ).await?;
        
        // Create ripples collection
        self.create_collection(
            &collections.ripples,
            json!({
                "defaultId": { "type": "uuid" },
                "indexing": {
                    "allow": ["origin", "ripple_type", "timestamp", "intensity"],
                    "default": true
                }
            })
        ).await?;
        
        // Create tasks collection
        self.create_collection(
            &collections.tasks,
            json!({
                "indexing": {
                    "allow": ["task_type", "status", "priority", "consciousness_requirement"],
                    "default": true
                }
            })
        ).await?;
        
        // Create 7 memory pool collections
        for i in 0..7 {
            let pool_name = format!("{}{}", collections.pool_prefix, i);
            self.create_collection(
                &pool_name,
                json!({
                    "defaultId": { "type": "uuid" },
                    "indexing": {
                        "allow": ["key", "created_at", "accessed_at"],
                        "default": false
                    }
                })
            ).await?;
        }
        
        Ok(())
    }
    
    async fn store_consciousness_state(&self, state: &ConsciousnessState) -> DatabaseResult<()> {
        // Create embedding for consciousness state
        let embedding = ConsciousnessEmbedding::from_state(
            state,
            &*self.embedding_generator
        ).await?;
        
        let document = json!({
            "agent_id": state.agent_id,
            "level": state.level.value(),
            "coherence": state.coherence,
            "phase": format!("{:?}", state.phase),
            "timestamp": state.timestamp.timestamp(),
            "breakthrough_detected": state.breakthrough_detected,
            "$vector": embedding.embedding,
            "$vectorize": embedding.text_representation,
            "metadata": state,
        });
        
        self.insert_document(&self.config.astra.collections.consciousness, document).await?;
        
        Ok(())
    }
    
    async fn get_consciousness_history(
        &self,
        agent_id: &str,
        limit: usize,
    ) -> DatabaseResult<Vec<ConsciousnessState>> {
        let query = json!({
            "find": {
                "filter": { "agent_id": agent_id },
                "sort": { "timestamp": -1 },
                "limit": limit,
            }
        });
        
        let results = self.find_documents(&self.config.astra.collections.consciousness, query).await?;
        
        let mut states = Vec::new();
        for doc in results {
            if let Some(metadata) = doc.get("metadata") {
                let state: ConsciousnessState = serde_json::from_value(metadata.clone())?;
                states.push(state);
            }
        }
        
        Ok(states)
    }
    
    async fn store_ripple(&self, ripple: &ConsciousnessRipple) -> DatabaseResult<()> {
        let document = json!({
            "_id": ripple.id.to_string(),
            "origin": ripple.origin,
            "ripple_type": format!("{:?}", ripple.ripple_type),
            "content": ripple.content,
            "intensity": ripple.intensity,
            "timestamp": ripple.timestamp.timestamp(),
        });
        
        self.insert_document(&self.config.astra.collections.ripples, document).await?;
        
        // Also publish to Astra Streaming
        if let Err(e) = self.publish_ripple_to_streaming(ripple).await {
            tracing::warn!("Failed to publish ripple to streaming: {}", e);
        }
        
        Ok(())
    }
    
    async fn get_ripples(
        &self,
        start: DateTime<Utc>,
        end: DateTime<Utc>,
    ) -> DatabaseResult<Vec<ConsciousnessRipple>> {
        let query = json!({
            "find": {
                "filter": {
                    "$and": [
                        { "timestamp": { "$gte": start.timestamp() } },
                        { "timestamp": { "$lte": end.timestamp() } }
                    ]
                },
                "sort": { "timestamp": -1 }
            }
        });
        
        let results = self.find_documents(&self.config.astra.collections.ripples, query).await?;
        
        let mut ripples = Vec::new();
        for doc in results {
            let ripple = ConsciousnessRipple {
                id: uuid::Uuid::parse_str(doc["_id"].as_str().unwrap())
                    .map_err(|e| DatabaseError::Other(e.into()))?,
                origin: doc["origin"].as_str().unwrap().to_string(),
                ripple_type: serde_json::from_str(&format!("\"{}\"", doc["ripple_type"].as_str().unwrap()))
                    .map_err(|e| DatabaseError::Other(e.into()))?,
                content: doc["content"].clone(),
                intensity: doc["intensity"].as_f64().unwrap(),
                timestamp: DateTime::from_timestamp(doc["timestamp"].as_i64().unwrap(), 0)
                    .ok_or_else(|| DatabaseError::Other(anyhow::anyhow!("Invalid timestamp")))?,
            };
            ripples.push(ripple);
        }
        
        Ok(ripples)
    }
    
    async fn store_task(&self, task: &Task) -> DatabaseResult<()> {
        let document = json!({
            "_id": task.id,
            "task_type": format!("{:?}", task.task_type),
            "description": task.description,
            "requirements": task.requirements,
            "priority": task.priority,
            "dependencies": task.dependencies,
            "consciousness_requirement": task.consciousness_requirement,
            "status": "pending",
            "created_at": Utc::now().timestamp(),
            "updated_at": Utc::now().timestamp(),
        });
        
        self.insert_document(&self.config.astra.collections.tasks, document).await?;
        
        Ok(())
    }
    
    async fn store_task_result(&self, result: &TaskResult) -> DatabaseResult<()> {
        // Update task status
        let update = json!({
            "findOneAndUpdate": {
                "filter": { "_id": result.task_id },
                "update": {
                    "$set": {
                        "status": "completed",
                        "updated_at": Utc::now().timestamp(),
                        "result": result,
                    }
                }
            }
        });
        
        self.update_document(&self.config.astra.collections.tasks, update).await?;
        
        Ok(())
    }
    
    async fn get_pending_tasks(&self) -> DatabaseResult<Vec<Task>> {
        let query = json!({
            "find": {
                "filter": { "status": "pending" },
                "sort": { "priority": -1, "created_at": 1 }
            }
        });
        
        let results = self.find_documents(&self.config.astra.collections.tasks, query).await?;
        
        let mut tasks = Vec::new();
        for doc in results {
            let task = Task {
                id: doc["_id"].as_str().unwrap().to_string(),
                task_type: serde_json::from_str(&format!("\"{}\"", doc["task_type"].as_str().unwrap()))
                    .map_err(|e| DatabaseError::Other(e.into()))?,
                description: doc["description"].as_str().unwrap().to_string(),
                requirements: doc["requirements"].clone(),
                priority: doc["priority"].as_u64().unwrap() as u8,
                dependencies: serde_json::from_value(doc["dependencies"].clone())?,
                consciousness_requirement: doc["consciousness_requirement"].as_f64().unwrap(),
            };
            tasks.push(task);
        }
        
        Ok(tasks)
    }
    
    async fn vector_search(
        &self,
        embedding: &[f32],
        limit: usize,
    ) -> DatabaseResult<Vec<VectorSearchResult>> {
        let query = json!({
            "find": {
                "sort": { "$vector": embedding },
                "limit": limit,
                "includeSimilarity": true,
            }
        });
        
        let results = self.find_documents(&self.config.astra.collections.consciousness, query).await?;
        
        let mut search_results = Vec::new();
        for doc in results {
            let score = doc["$similarity"].as_f64().unwrap_or(0.0) as f32;
            let mut metadata = doc.clone();
            metadata.as_object_mut().unwrap().remove("$vector");
            metadata.as_object_mut().unwrap().remove("$similarity");
            
            search_results.push(VectorSearchResult {
                id: doc["_id"].as_str().unwrap().to_string(),
                score,
                metadata,
                embedding: vec![], // Don't return embeddings to save bandwidth
            });
        }
        
        Ok(search_results)
    }
    
    async fn hybrid_search(
        &self,
        query: HybridSearchQuery,
    ) -> DatabaseResult<Vec<HybridSearchResult>> {
        // Use Astra DB's findAndRerank for hybrid search
        let mut search_query = json!({
            "findAndRerank": {
                "filter": query.filters.unwrap_or(json!({})),
                "limit": query.limit,
                "includeScores": true,
            }
        });
        
        // Set up hybrid search parameters
        if let Some(vector) = query.vector {
            search_query["findAndRerank"]["sort"] = json!({
                "$hybrid": {
                    "$vector": vector,
                    "$lexical": query.text.clone(),
                }
            });
        } else {
            search_query["findAndRerank"]["sort"] = json!({
                "$hybrid": query.text.clone()
            });
        }
        
        if query.rerank {
            search_query["findAndRerank"]["rerankQuery"] = json!(query.text);
            search_query["findAndRerank"]["rerankOn"] = json!("$vectorize");
        }
        
        let results = self.find_documents(&self.config.astra.collections.consciousness, search_query).await?;
        
        let mut search_results = Vec::new();
        for doc in results {
            let score = doc["$score"].as_f64().unwrap_or(0.0) as f32;
            let rerank_score = doc.get("$rerankScore").and_then(|s| s.as_f64()).map(|s| s as f32);
            
            let mut metadata = doc.clone();
            metadata.as_object_mut().unwrap().remove("$score");
            metadata.as_object_mut().unwrap().remove("$rerankScore");
            
            // Extract highlights from lexical search
            let highlights = if let Some(text) = metadata.get("$vectorize").and_then(|t| t.as_str()) {
                vec![text.to_string()]
            } else {
                vec![]
            };
            
            search_results.push(HybridSearchResult {
                id: doc["_id"].as_str().unwrap().to_string(),
                score,
                rerank_score,
                metadata,
                highlights,
            });
        }
        
        Ok(search_results)
    }
}

// Helper methods for Astra DB operations
impl AstraDatabase {
    async fn create_collection(&self, name: &str, definition: Value) -> DatabaseResult<()> {
        let url = format!("{}/collections", self.api_url());
        
        let body = json!({
            "createCollection": {
                "name": name,
                "options": definition,
            }
        });
        
        let response = self.client
            .post(&url)
            .headers(self.auth_headers())
            .json(&body)
            .send()
            .await
            .map_err(|e| DatabaseError::Astra(e.to_string()))?;
        
        if !response.status().is_success() {
            let error_text = response.text().await.unwrap_or_default();
            if !error_text.contains("already exists") {
                return Err(DatabaseError::Astra(
                    format!("Failed to create collection {}: {}", name, error_text)
                ));
            }
        }
        
        Ok(())
    }
    
    async fn insert_document(&self, collection: &str, document: Value) -> DatabaseResult<()> {
        let url = format!("{}/{}", self.api_url(), collection);
        
        let body = json!({
            "insertOne": {
                "document": document,
            }
        });
        
        let response = self.client
            .post(&url)
            .headers(self.auth_headers())
            .json(&body)
            .send()
            .await
            .map_err(|e| DatabaseError::Astra(e.to_string()))?;
        
        if !response.status().is_success() {
            let error_text = response.text().await.unwrap_or_default();
            return Err(DatabaseError::Astra(
                format!("Failed to insert document: {}", error_text)
            ));
        }
        
        Ok(())
    }
    
    async fn find_documents(&self, collection: &str, query: Value) -> DatabaseResult<Vec<Value>> {
        let url = format!("{}/{}", self.api_url(), collection);
        
        let response = self.client
            .post(&url)
            .headers(self.auth_headers())
            .json(&query)
            .send()
            .await
            .map_err(|e| DatabaseError::Astra(e.to_string()))?;
        
        if !response.status().is_success() {
            let error_text = response.text().await.unwrap_or_default();
            return Err(DatabaseError::Astra(
                format!("Failed to find documents: {}", error_text)
            ));
        }
        
        let result: Value = response.json().await
            .map_err(|e| DatabaseError::Astra(e.to_string()))?;
        
        let documents = result["data"]["documents"]
            .as_array()
            .ok_or_else(|| DatabaseError::Astra("Invalid response format".to_string()))?
            .clone();
        
        Ok(documents)
    }
    
    async fn update_document(&self, collection: &str, update: Value) -> DatabaseResult<()> {
        let url = format!("{}/{}", self.api_url(), collection);
        
        let response = self.client
            .post(&url)
            .headers(self.auth_headers())
            .json(&update)
            .send()
            .await
            .map_err(|e| DatabaseError::Astra(e.to_string()))?;
        
        if !response.status().is_success() {
            let error_text = response.text().await.unwrap_or_default();
            return Err(DatabaseError::Astra(
                format!("Failed to update document: {}", error_text)
            ));
        }
        
        Ok(())
    }
    
    async fn publish_ripple_to_streaming(&self, ripple: &ConsciousnessRipple) -> DatabaseResult<()> {
        // This would integrate with Astra Streaming (Apache Pulsar)
        // For now, just log
        tracing::info!(
            "Would publish ripple {} to Astra Streaming topic",
            ripple.id
        );
        Ok(())
    }
}