use crate::models::{Chunk, Document, Node};
use async_trait::async_trait;
use thiserror::Error;

#[derive(Error, Debug)]
pub enum CerebroError {
    #[error("Storage error: {0}")]
    StorageError(String),
    #[error("Embedding error: {0}")]
    EmbeddingError(String),
    #[error("Ingestion error: {0}")]
    IngestionError(String),
}

pub type Result<T> = std::result::Result<T, CerebroError>;

/// Handles converting raw bytes or external systems into standardized Documents.
#[async_trait]
pub trait Ingestor: Send + Sync {
    async fn ingest(&self, source: &str) -> Result<Vec<Document>>;
}

/// Chunks a large document into pieces suitable for embedding.
pub trait Chunker: Send + Sync {
    fn chunk(&self, document: &Document) -> Result<Vec<Chunk>>;
}

/// Generates vector embeddings for strings.
#[async_trait]
pub trait Embedder: Send + Sync {
    async fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>>;

    async fn embed_query(&self, text: &str) -> Result<Vec<f32>> {
        let res = self.embed(&[text]).await?;
        Ok(res.into_iter().next().unwrap_or_default())
    }
}

/// The core target for scalable vector storage.
#[async_trait]
pub trait VectorStore: Send + Sync {
    async fn upsert(&self, nodes: Vec<Node>) -> Result<()>;
    async fn get(&self, node_ids: &[&str]) -> Result<Vec<Node>>;
    async fn search(&self, text_query: &str, embedding: &[f32], top_k: usize) -> Result<Vec<(Node, f32)>>;
    async fn delete_document(&self, doc_id: &str) -> Result<()>;
    async fn get_all_nodes(&self) -> Result<Vec<Node>>;
}

/// Key-Value store trait used for Working Memory (fast state retrieval).
#[async_trait]
pub trait KVStore: Send + Sync {
    async fn set(&self, key: &str, value: &str) -> Result<()>;
    async fn get(&self, key: &str) -> Result<Option<String>>;
}
