// RLM Storage Layer
// Provides persistence for chunks, buffers, and execution history

pub mod surrealdb;

// Re-export main storage type
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
pub use surrealdb::SurrealDBStorage;

/// A chunk from a document
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Chunk {
    pub chunk_id: String,
    pub doc_id: String,
    pub content: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub embedding: Option<Vec<f32>>,
    pub start_idx: usize,
    pub end_idx: usize,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub metadata: Option<serde_json::Value>,
    pub created_at: String,
}

/// A buffer for pass-by-reference large contexts
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Buffer {
    pub buffer_id: String,
    pub content: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub metadata: Option<serde_json::Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub expires_at: Option<String>,
    pub created_at: String,
}

/// Execution history record
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionHistory {
    pub execution_id: String,
    pub doc_id: String,
    pub query: String,
    pub chunks_used: Vec<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<String>,
    pub duration_ms: u64,
    pub cost_cents: f64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub provider: Option<String>,
    pub success: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error_message: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub metadata: Option<serde_json::Value>,
    pub created_at: String,
    pub executed_at: String,
}

/// Storage trait for RLM operations
#[async_trait]
pub trait Storage: Send + Sync {
    /// Save a chunk to storage
    async fn save_chunk(&self, chunk: Chunk) -> crate::Result<()>;

    /// Get chunks by document ID
    async fn get_chunks(&self, doc_id: &str) -> crate::Result<Vec<Chunk>>;

    /// Get a specific chunk by ID
    async fn get_chunk(&self, chunk_id: &str) -> crate::Result<Option<Chunk>>;

    /// Search chunks by embedding similarity
    async fn search_by_embedding(
        &self,
        embedding: &[f32],
        limit: usize,
    ) -> crate::Result<Vec<Chunk>>;

    /// Save a buffer to storage
    async fn save_buffer(&self, buffer: Buffer) -> crate::Result<()>;

    /// Get a buffer by ID
    async fn get_buffer(&self, buffer_id: &str) -> crate::Result<Option<Buffer>>;

    /// Delete expired buffers
    async fn cleanup_expired_buffers(&self) -> crate::Result<u64>;

    /// Save execution history
    async fn save_execution(&self, execution: ExecutionHistory) -> crate::Result<()>;

    /// Get execution history by document ID
    async fn get_executions(
        &self,
        doc_id: &str,
        limit: usize,
    ) -> crate::Result<Vec<ExecutionHistory>>;

    /// Delete chunks by document ID
    async fn delete_chunks(&self, doc_id: &str) -> crate::Result<u64>;
}

// Mock storage for testing
#[cfg(test)]
pub use mock::MockStorage;

#[cfg(test)]
mod mock {
    use std::collections::HashMap;
    use std::sync::{Arc, Mutex};

    use super::*;

    /// Mock storage implementation for testing
    pub struct MockStorage {
        pub(crate) chunks: Arc<Mutex<HashMap<String, Vec<Chunk>>>>,
    }

    impl Default for MockStorage {
        fn default() -> Self {
            Self::new()
        }
    }

    impl MockStorage {
        pub fn new() -> Self {
            Self {
                chunks: Arc::new(Mutex::new(HashMap::new())),
            }
        }
    }

    #[async_trait]
    impl Storage for MockStorage {
        async fn save_chunk(&self, chunk: Chunk) -> crate::Result<()> {
            let mut chunks = self.chunks.lock().unwrap();
            chunks.entry(chunk.doc_id.clone()).or_default().push(chunk);
            Ok(())
        }

        async fn get_chunks(&self, doc_id: &str) -> crate::Result<Vec<Chunk>> {
            let chunks = self.chunks.lock().unwrap();
            Ok(chunks.get(doc_id).cloned().unwrap_or_default())
        }

        async fn get_chunk(&self, chunk_id: &str) -> crate::Result<Option<Chunk>> {
            let chunks = self.chunks.lock().unwrap();
            for doc_chunks in chunks.values() {
                if let Some(chunk) = doc_chunks.iter().find(|c| c.chunk_id == chunk_id) {
                    return Ok(Some(chunk.clone()));
                }
            }
            Ok(None)
        }

        async fn search_by_embedding(
            &self,
            _embedding: &[f32],
            limit: usize,
        ) -> crate::Result<Vec<Chunk>> {
            let chunks = self.chunks.lock().unwrap();
            let all_chunks: Vec<Chunk> = chunks.values().flatten().take(limit).cloned().collect();
            Ok(all_chunks)
        }

        async fn save_buffer(&self, _buffer: Buffer) -> crate::Result<()> {
            Ok(())
        }

        async fn get_buffer(&self, _buffer_id: &str) -> crate::Result<Option<Buffer>> {
            Ok(None)
        }

        async fn cleanup_expired_buffers(&self) -> crate::Result<u64> {
            Ok(0)
        }

        async fn save_execution(&self, _execution: ExecutionHistory) -> crate::Result<()> {
            Ok(())
        }

        async fn get_executions(
            &self,
            _doc_id: &str,
            _limit: usize,
        ) -> crate::Result<Vec<ExecutionHistory>> {
            Ok(vec![])
        }

        async fn delete_chunks(&self, doc_id: &str) -> crate::Result<u64> {
            let mut chunks = self.chunks.lock().unwrap();
            if let Some(doc_chunks) = chunks.remove(doc_id) {
                Ok(doc_chunks.len() as u64)
            } else {
                Ok(0)
            }
        }
    }
}
