//! Vector store for RAG and semantic search

use anyhow::Result;
use rusqlite::{Connection, params};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

/// Document for vector store
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Document {
    pub id: String,
    pub content: String,
    pub metadata: serde_json::Value,
    pub embedding: Vec<f32>,
}

/// Vector store using SQLite with simple embeddings
pub struct VectorStore {
    conn: Connection,
}

impl VectorStore {
    /// Create new vector store
    pub async fn new(db_path: &str) -> Result<Self> {
        let conn = Connection::open(db_path)?;
        
        // Create tables
        conn.execute(
            "CREATE TABLE IF NOT EXISTS documents (
                id TEXT PRIMARY KEY,
                content TEXT NOT NULL,
                metadata TEXT,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP
            )",
            [],
        )?;
        
        conn.execute(
            "CREATE TABLE IF NOT EXISTS embeddings (
                document_id TEXT PRIMARY KEY,
                embedding BLOB NOT NULL,
                FOREIGN KEY (document_id) REFERENCES documents(id)
            )",
            [],
        )?;
        
        Ok(Self { conn })
    }
    
    /// Add document to store
    pub fn add_document(&self, doc: &Document) -> Result<()> {
        let tx = self.conn.transaction()?;
        
        // Insert document
        tx.execute(
            "INSERT OR REPLACE INTO documents (id, content, metadata) VALUES (?1, ?2, ?3)",
            params![doc.id, doc.content, serde_json::to_string(&doc.metadata)?],
        )?;
        
        // Insert embedding (store as blob)
        let embedding_bytes: Vec<u8> = doc.embedding
            .iter()
            .flat_map(|&f| f.to_le_bytes().to_vec())
            .collect();
        
        tx.execute(
            "INSERT OR REPLACE INTO embeddings (document_id, embedding) VALUES (?1, ?2)",
            params![doc.id, embedding_bytes],
        )?;
        
        tx.commit()?;
        
        Ok(())
    }
    
    /// Search documents by similarity (simple cosine similarity)
    pub fn search(&self, query_embedding: &[f32], limit: usize) -> Result<Vec<Document>> {
        let mut stmt = self.conn.prepare(
            "SELECT d.id, d.content, d.metadata, e.embedding 
             FROM documents d 
             JOIN embeddings e ON d.id = e.document_id 
             ORDER BY d.created_at DESC 
             LIMIT ?1"
        )?;
        
        let docs = stmt.query_map(params![limit], |row| {
            let id: String = row.get(0)?;
            let content: String = row.get(1)?;
            let metadata: String = row.get(2)?;
            let embedding_blob: Vec<u8> = row.get(3)?;
            
            // Convert blob back to f32 vector
            let embedding: Vec<f32> = embedding_blob
                .chunks(4)
                .map(|chunk| {
                    let bytes: [u8; 4] = chunk.try_into().unwrap_or([0; 4]);
                    f32::from_le_bytes(bytes)
                })
                .collect();
            
            Ok(Document {
                id,
                content,
                metadata: serde_json::from_str(&metadata).unwrap_or_default(),
                embedding,
            })
        })?;
        
        let mut results: Vec<Document> = docs.filter_map(|r| r.ok()).collect();
        
        // Sort by cosine similarity
        results.sort_by(|a, b| {
            let sim_a = cosine_similarity(query_embedding, &a.embedding);
            let sim_b = cosine_similarity(query_embedding, &b.embedding);
            sim_b.partial_cmp(&sim_a).unwrap_or(std::cmp::Ordering::Equal)
        });
        
        Ok(results)
    }
    
    /// Get document by ID
    pub fn get_document(&self, id: &str) -> Result<Option<Document>> {
        let mut stmt = self.conn.prepare(
            "SELECT d.id, d.content, d.metadata, e.embedding 
             FROM documents d 
             JOIN embeddings e ON d.id = e.document_id 
             WHERE d.id = ?1"
        )?;
        
        let doc = stmt.query_row(params![id], |row| {
            let id: String = row.get(0)?;
            let content: String = row.get(1)?;
            let metadata: String = row.get(2)?;
            let embedding_blob: Vec<u8> = row.get(3)?;
            
            let embedding: Vec<f32> = embedding_blob
                .chunks(4)
                .map(|chunk| {
                    let bytes: [u8; 4] = chunk.try_into().unwrap_or([0; 4]);
                    f32::from_le_bytes(bytes)
                })
                .collect();
            
            Ok(Document {
                id,
                content,
                metadata: serde_json::from_str(&metadata).unwrap_or_default(),
                embedding,
            })
        }).ok();
        
        Ok(doc.ok().flatten())
    }
    
    /// Delete document
    pub fn delete_document(&self, id: &str) -> Result<()> {
        self.conn.execute(
            "DELETE FROM documents WHERE id = ?1",
            params![id],
        )?;
        
        self.conn.execute(
            "DELETE FROM embeddings WHERE document_id = ?1",
            params![id],
        )?;
        
        Ok(())
    }
    
    /// Get document count
    pub fn count(&self) -> Result<usize> {
        let count: usize = self.conn.query_row(
            "SELECT COUNT(*) FROM documents",
            [],
            |row| row.get(0),
        )?;
        
        Ok(count)
    }
}

/// Calculate cosine similarity between two vectors
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    let dot: 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 / (norm_a * norm_b)
    }
}

/// Simple embedding function (placeholder - in production, use real embeddings)
pub fn simple_embed(text: &str) -> Vec<f32> {
    // Simple hash-based embedding (NOT production quality!)
    // In production, use a real embedding model via Rig
    let mut embedding = vec![0.0; 384]; // 384-dimensional embedding
    
    for (i, byte) in text.bytes().enumerate() {
        embedding[i % 384] += byte as f32 / 255.0;
    }
    
    // Normalize
    let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > 0.0 {
        for x in embedding.iter_mut() {
            *x /= norm;
        }
    }
    
    embedding
}
