use crate::query_optimizer::{QueryContext, QueryOptimizer, RRFConfig, RRFFusion, SearchStrategy};
use crate::opensearch_client::OpenSearchClient;
use anyhow::Result;
use mem_llm::EmbeddingsClient;
use mem_store::VectorStore;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::Instant;

/// Hybrid Query Result with score breakdown
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct HybridQueryResult {
    pub id: String,
    pub text: String,
    pub source: String,
    pub level: String,
    pub breadcrumb: Vec<String>,
    
    // Scoring breakdown
    pub final_score: f32,
    pub semantic_score: Option<f32>,    // From pgvector
    pub lexical_score: Option<f32>,     // From OpenSearch
    pub fusion_method: String,          // "rrf" or "weighted_linear"
    pub rank: usize,
    pub retrieval_engine: String,       // "semantic_only", "lexical_only", or "hybrid"
}

/// Hybrid Query Response
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct HybridQueryResponse {
    pub query: String,
    pub project: String,
    pub search_strategy: String,
    pub strategy_confidence: f32,
    pub results: Vec<HybridQueryResult>,
    pub metrics: QueryMetrics,
}

/// Query execution metrics
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct QueryMetrics {
    pub total_time_ms: u128,
    pub semantic_time_ms: Option<u128>,
    pub lexical_time_ms: Option<u128>,
    pub fusion_time_ms: u128,
    pub semantic_results_count: Option<usize>,
    pub lexical_results_count: Option<usize>,
    pub final_results_count: usize,
}

/// Hybrid Query Worker: orchestrates parallel retrieval
pub struct HybridQueryWorker {
    optimizer: Arc<QueryOptimizer>,
    vector_store: Arc<VectorStore>,
    embeddings: Arc<EmbeddingsClient>,
    opensearch: Option<Arc<OpenSearchClient>>,
    rrf_config: RRFConfig,
}

impl HybridQueryWorker {
    pub fn new(
        vector_store: Arc<VectorStore>,
        embeddings: Arc<EmbeddingsClient>,
        opensearch: Option<Arc<OpenSearchClient>>,
    ) -> Self {
        Self {
            optimizer: Arc::new(QueryOptimizer::new()),
            vector_store,
            embeddings,
            opensearch,
            rrf_config: RRFConfig::default(),
        }
    }

    /// Main entry point: hybrid query with full orchestration
    pub async fn query(
        &self,
        project: &str,
        question: &str,
        limit: i64,
        jwt_token: &str,
    ) -> Result<HybridQueryResponse> {
        let start = Instant::now();

        // Stage 1: Optimize query
        let mut query_ctx = self.optimizer.optimize_query(question).await?;
        
        // Stage 2: Generate embedding
        query_ctx.embedding = Some(self.embeddings.embed_one(question).await?);
        
        // Stage 3: Execute retrieval based on strategy
        let (semantic_results, lexical_results, metrics) = match &query_ctx.search_strategy {
            SearchStrategy::Hybrid => {
                self.retrieve_hybrid(
                    project,
                    &query_ctx,
                    limit,
                    jwt_token,
                )
                .await?
            }
            SearchStrategy::SemanticOnly => {
                let sem_results = self.retrieve_semantic(project, &query_ctx, limit).await?;
                (Some(sem_results), None, QueryMetrics::default())
            }
            SearchStrategy::LexicalOnly => {
                let lex_results = self.retrieve_lexical(project, &query_ctx, limit, jwt_token).await?;
                (None, Some(lex_results), QueryMetrics::default())
            }
            SearchStrategy::LexicalFirst => {
                self.retrieve_cascading(
                    project,
                    &query_ctx,
                    limit,
                    jwt_token,
                )
                .await?
            }
        };

        // Stage 4: Fuse results
        let fusion_start = Instant::now();
        let fused = self.fuse_results(semantic_results, lexical_results)?;
        let fusion_time_ms = fusion_start.elapsed().as_millis();

        // Stage 5: Build response
        let results = self.build_results(fused, &query_ctx).await?;

        let mut metrics = metrics;
        metrics.total_time_ms = start.elapsed().as_millis();
        metrics.fusion_time_ms = fusion_time_ms;
        metrics.final_results_count = results.len();

        Ok(HybridQueryResponse {
            query: question.to_string(),
            project: project.to_string(),
            search_strategy: format!("{:?}", query_ctx.search_strategy),
            strategy_confidence: query_ctx.confidence,
            results,
            metrics,
        })
    }

    /// Hybrid retrieval: parallel pgvector + OpenSearch
    async fn retrieve_hybrid(
        &self,
        project: &str,
        query_ctx: &QueryContext,
        limit: i64,
        jwt_token: &str,
    ) -> Result<(Option<Vec<(String, f32)>>, Option<Vec<(String, f32)>>, QueryMetrics)> {
        let embedding = query_ctx
            .embedding
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("no embedding generated"))?;

        // Parallel execution
        let semantic_fut = self.retrieve_semantic(project, query_ctx, limit);
        let lexical_fut = self.retrieve_lexical(project, query_ctx, limit, jwt_token);

        let sem_start = Instant::now();
        let (semantic_results, lexical_results) = tokio::try_join!(semantic_fut, lexical_fut)?;
        let sem_time = sem_start.elapsed().as_millis();

        let metrics = QueryMetrics {
            semantic_time_ms: Some(sem_time),
            lexical_time_ms: Some(sem_time),  // Parallel, so roughly same
            semantic_results_count: Some(semantic_results.len()),
            lexical_results_count: Some(lexical_results.len()),
            ..Default::default()
        };

        Ok((Some(semantic_results), Some(lexical_results), metrics))
    }

    /// Cascading retrieval: lexical → semantic
    async fn retrieve_cascading(
        &self,
        project: &str,
        query_ctx: &QueryContext,
        limit: i64,
        jwt_token: &str,
    ) -> Result<(Option<Vec<(String, f32)>>, Option<Vec<(String, f32)>>, QueryMetrics)> {
        // Stage 1: Lexical search (narrow down)
        let lex_start = Instant::now();
        let lexical_results = self.retrieve_lexical(project, query_ctx, limit * 4, jwt_token).await?;
        let lex_time = lex_start.elapsed().as_millis();

        // Extract chunk IDs from lexical results
        let chunk_ids: Vec<String> = lexical_results.iter().map(|(id, _)| id.clone()).collect();

        // Stage 2: Semantic rerank (on narrowed set)
        let sem_start = Instant::now();
        let semantic_results = self
            .retrieve_semantic_with_ids(project, query_ctx, limit, &chunk_ids)
            .await?;
        let sem_time = sem_start.elapsed().as_millis();

        let metrics = QueryMetrics {
            lexical_time_ms: Some(lex_time),
            semantic_time_ms: Some(sem_time),
            lexical_results_count: Some(lexical_results.len()),
            semantic_results_count: Some(semantic_results.len()),
            ..Default::default()
        };

        Ok((Some(semantic_results), Some(lexical_results), metrics))
    }

    /// Retrieve from pgvector (semantic search)
    async fn retrieve_semantic(
        &self,
        project: &str,
        query_ctx: &QueryContext,
        limit: i64,
    ) -> Result<Vec<(String, f32)>> {
        let embedding = query_ctx
            .embedding
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("no embedding generated"))?;

        // Query pgvector with filters
        let results = self
            .vector_store
            .search(embedding, project, limit, None)
            .await?;

        // Convert to (id, score) tuples
        let scored: Vec<(String, f32)> = results
            .into_iter()
            .map(|(id, score, _)| (id, score))
            .collect();

        Ok(scored)
    }

    /// Retrieve from pgvector with specific chunk IDs (for cascading)
    async fn retrieve_semantic_with_ids(
        &self,
        project: &str,
        query_ctx: &QueryContext,
        limit: i64,
        chunk_ids: &[String],
    ) -> Result<Vec<(String, f32)>> {
        let embedding = query_ctx
            .embedding
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("no embedding generated"))?;

        // Query pgvector filtered by chunk IDs
        let results = self
            .vector_store
            .search_with_ids(embedding, project, limit, chunk_ids)
            .await?;

        let scored: Vec<(String, f32)> = results
            .into_iter()
            .map(|(id, score, _)| (id, score))
            .collect();

        Ok(scored)
    }

    /// Retrieve from OpenSearch (lexical search)
    async fn retrieve_lexical(
        &self,
        project: &str,
        query_ctx: &QueryContext,
        limit: i64,
        jwt_token: &str,
    ) -> Result<Vec<(String, f32)>> {
        let opensearch = self
            .opensearch
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("OpenSearch not configured"))?;

        // Query OpenSearch with JWT auth
        let results = opensearch
            .lexical_search(&query_ctx.normalized_query, limit as usize, jwt_token)
            .await?;

        // Convert to (id, score) tuples
        let scored: Vec<(String, f32)> = results
            .into_iter()
            .map(|(id, score, _, _, _)| (id, score))
            .collect();

        Ok(scored)
    }

    /// Fuse semantic and lexical results using RRF
    fn fuse_results(
        &self,
        semantic: Option<Vec<(String, f32)>>,
        lexical: Option<Vec<(String, f32)>>,
    ) -> Result<Vec<(String, f32)>> {
        match (semantic, lexical) {
            (Some(sem), Some(lex)) => {
                // Use RRF for fusion
                let fusion = RRFFusion::new(self.rrf_config.clone());
                Ok(fusion.fuse(sem, lex))
            }
            (Some(sem), None) => {
                // Semantic only: return top-k
                let mut results = sem;
                results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
                results.truncate(self.rrf_config.final_k);
                Ok(results)
            }
            (None, Some(lex)) => {
                // Lexical only: return top-k
                let mut results = lex;
                results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
                results.truncate(self.rrf_config.final_k);
                Ok(results)
            }
            (None, None) => Err(anyhow::anyhow!("no results from either engine")),
        }
    }

    /// Build response with enriched metadata
    async fn build_results(
        &self,
        fused: Vec<(String, f32)>,
        query_ctx: &QueryContext,
    ) -> Result<Vec<HybridQueryResult>> {
        let mut results = Vec::new();

        for (rank, (id, score)) in fused.into_iter().enumerate() {
            // Fetch full chunk metadata from database
            let chunk = self.vector_store.get_chunk(&id).await?;

            results.push(HybridQueryResult {
                id: id.clone(),
                text: chunk.text,
                source: chunk.source,
                level: chunk.level.unwrap_or_default(),
                breadcrumb: chunk.breadcrumb.unwrap_or_default(),
                final_score: score,
                semantic_score: None,    // Would need to track separately
                lexical_score: None,      // Would need to track separately
                fusion_method: "rrf".to_string(),
                rank: rank + 1,
                retrieval_engine: format!("{:?}", query_ctx.search_strategy),
            });
        }

        Ok(results)
    }
}

impl Default for QueryMetrics {
    fn default() -> Self {
        Self {
            total_time_ms: 0,
            semantic_time_ms: None,
            lexical_time_ms: None,
            fusion_time_ms: 0,
            semantic_results_count: None,
            lexical_results_count: None,
            final_results_count: 0,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // These tests require mock implementations of VectorStore and EmbeddingsClient
    // Placeholder tests for structure verification

    #[test]
    fn test_hybrid_response_structure() {
        let resp = HybridQueryResponse {
            query: "test".to_string(),
            project: "poimen".to_string(),
            search_strategy: "Hybrid".to_string(),
            strategy_confidence: 0.95,
            results: vec![],
            metrics: QueryMetrics::default(),
        };

        assert_eq!(resp.query, "test");
        assert_eq!(resp.strategy_confidence, 0.95);
    }
}