use anyhow::Result;
use mem_llm::{EmbeddingsClient, RerankClient};
use mem_store::VectorStore;
use pgvector::Vector;
use serde::{Deserialize, Serialize};

/// Query result with provenance
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryResult {
    pub level: String,  // "L0", "L1", "L2", "corpus"
    pub score: f32,
    pub text: String,
    pub source: Option<String>,
    pub provenance: Vec<String>,  // parent IDs
}

/// Query worker — semantic search + reranking
pub struct QueryWorker {
    vector_store: std::sync::Arc<VectorStore>,
    embeddings: std::sync::Arc<EmbeddingsClient>,
    reranker: std::sync::Arc<RerankClient>,
}

impl QueryWorker {
    /// Create query worker
    pub fn new(
        vector_store: VectorStore,
        embeddings: EmbeddingsClient,
        reranker: RerankClient,
    ) -> Self {
        Self {
            vector_store: std::sync::Arc::new(vector_store),
            embeddings: std::sync::Arc::new(embeddings),
            reranker: std::sync::Arc::new(reranker),
        }
    }

    /// Execute semantic query: embed -> search vector -> rerank -> result
    pub async fn query(
        &self,
        project: &str,
        question: &str,
        limit: Option<i64>,
    ) -> Result<Vec<QueryResult>> {
        let limit = limit.unwrap_or(5);

        // Embed the question
        let question_embedding = self.embeddings.embed(question).await?;

        // Search across all levels
        let mut candidates = Vec::new();

        // L2 synthesis (project-level)
        if let Some(l2_result) = self.vector_store.search_l2(project, &question_embedding).await? {
            candidates.push(QueryResult {
                level: "L2".to_string(),
                score: l2_result.score,
                text: l2_result.item.content.clone(),
                source: Some(format!("project:{}", project)),
                provenance: vec![l2_result.item.id.to_string()],
            });
        }

        // L1 per-query memories
        let l1_results = self.vector_store.search_l1(project, &question_embedding, limit).await?;
        for l1_result in l1_results {
            candidates.push(QueryResult {
                level: "L1".to_string(),
                score: l1_result.score,
                text: l1_result.item.content.clone(),
                source: Some(format!("query:{}", l1_result.item.query_id)),
                provenance: vec![l1_result.item.id.to_string()],
            });
        }

        // Reference corpus
        let corpus_results = self.vector_store.search_corpus(project, &question_embedding, limit).await?;
        for corpus_result in corpus_results {
            candidates.push(QueryResult {
                level: "corpus".to_string(),
                score: corpus_result.score,
                text: corpus_result.item.content.clone(),
                source: Some(format!("doc:{}", corpus_result.item.name)),
                provenance: vec![corpus_result.item.id.to_string()],
            });
        }

        // Rerank candidates by relevance to question
        // TODO: wire actual cross-encoder reranking
        // For now, return by vector similarity score
        candidates.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
        candidates.truncate(limit as usize);

        Ok(candidates)
    }

    /// Get project synthesis (L2) directly
    pub async fn get_synthesis(&self, project: &str) -> Result<Option<QueryResult>> {
        if let Some(l2) = self.vector_store.get_l2(project).await? {
            Ok(Some(QueryResult {
                level: "L2".to_string(),
                score: 1.0,
                text: l2.content,
                source: Some(format!("project:{}", project)),
                provenance: vec![l2.id.to_string()],
            }))
        } else {
            Ok(None)
        }
    }
}