"""
Enhanced Pipeline.
[AC-AISVC-RES-02] 增强策略 Pipeline，新端到端流程。
"""

import asyncio
import logging
import re
import time
from dataclasses import dataclass
from typing import Any

import numpy as np

from app.core.config import get_settings
from app.core.qdrant_client import QdrantClient, get_qdrant_client
from app.services.embedding.nomic_provider import NomicEmbeddingProvider
from app.services.retrieval.base import RetrievalHit, RetrievalResult
from app.services.retrieval.optimized_retriever import RRFCombiner
from app.services.retrieval.strategy.config import (
    HybridRetrievalConfig,
    PipelineConfig,
    RerankerConfig,
)
from app.services.retrieval.strategy.pipeline_base import (
    BasePipeline,
    PipelineContext,
    PipelineResult,
)

logger = logging.getLogger(__name__)
settings = get_settings()


@dataclass
class RetrievalCandidate:
    """检索候选结果。"""
    id: str
    text: str
    score: float
    vector_score: float = 0.0
    keyword_score: float = 0.0
    metadata: dict[str, Any] = None
    
    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}


class EnhancedPipeline(BasePipeline):
    """
    增强策略 Pipeline。【AC-AISVC-RES-02】
    
    新端到端流程：
    1. Dense 向量检索
    2. Keyword 关键词检索
    3. RRF 融合排序
    4. 可选重排
    """
    
    def __init__(
        self,
        config: PipelineConfig | None = None,
        reranker_config: RerankerConfig | None = None,
        qdrant_client: QdrantClient | None = None,
    ):
        self._config = config or PipelineConfig()
        self._reranker_config = reranker_config or RerankerConfig()
        self._qdrant_client = qdrant_client
        self._rrf_combiner = RRFCombiner(k=self._config.hybrid.rrf_k)
        self._embedding_provider: NomicEmbeddingProvider | None = None
    
    @property
    def name(self) -> str:
        return "enhanced_pipeline"
    
    @property
    def description(self) -> str:
        return "增强检索策略，支持 Dense + Keyword + RRF 组合检索。"
    
    async def _get_client(self) -> QdrantClient:
        if self._qdrant_client is None:
            self._qdrant_client = await get_qdrant_client()
        return self._qdrant_client
    
    async def _get_embedding_provider(self) -> NomicEmbeddingProvider:
        if self._embedding_provider is None:
            from app.services.embedding.factory import get_embedding_config_manager
            manager = get_embedding_config_manager()
            provider = await manager.get_provider()
            if isinstance(provider, NomicEmbeddingProvider):
                self._embedding_provider = provider
            else:
                self._embedding_provider = NomicEmbeddingProvider(
                    base_url=settings.ollama_base_url,
                    model=settings.ollama_embedding_model,
                    dimension=settings.qdrant_vector_size,
                )
        return self._embedding_provider
    
    async def retrieve(self, ctx: PipelineContext) -> PipelineResult:
        """执行增强检索流程。【AC-AISVC-RES-02】"""
        start_time = time.time()
        
        logger.info(
            f"[EnhancedPipeline] Starting retrieval: tenant={ctx.tenant_id}, "
            f"query={ctx.query[:50]}..."
        )
        
        try:
            provider = await self._get_embedding_provider()
            embedding_result = await provider.embed_query(ctx.query)
            
            candidates = await self._hybrid_retrieve(
                tenant_id=ctx.tenant_id,
                query=ctx.query,
                embedding_result=embedding_result,
                metadata_filter=ctx.metadata_filter.filter_dict if ctx.metadata_filter else None,
                kb_ids=ctx.kb_ids,
            )
            
            if self._reranker_config.enabled and ctx.use_reranker:
                candidates = await self._rerank(
                    candidates=candidates,
                    query=ctx.query,
                )
            
            top_k = self._config.top_k
            final_candidates = candidates[:top_k]
            
            hits = [
                RetrievalHit(
                    text=c.text,
                    score=c.score,
                    source=self.name,
                    metadata=c.metadata,
                )
                for c in final_candidates
                if c.score >= self._config.score_threshold
            ]
            
            latency_ms = (time.time() - start_time) * 1000
            
            logger.info(
                f"[EnhancedPipeline] Retrieval completed: hits={len(hits)}, "
                f"latency_ms={latency_ms:.2f}"
            )
            
            result = RetrievalResult(
                hits=hits,
                diagnostics={
                    "total_candidates": len(candidates),
                    "after_rerank": self._reranker_config.enabled and ctx.use_reranker,
                },
            )
            
            return PipelineResult(
                retrieval_result=result,
                pipeline_name=self.name,
                used_reranker=self._reranker_config.enabled and ctx.use_reranker,
                metadata_filter_applied=ctx.metadata_filter is not None,
                latency_ms=latency_ms,
                diagnostics={
                    "dense_weight": self._config.hybrid.dense_weight,
                    "keyword_weight": self._config.hybrid.keyword_weight,
                },
            )
            
        except Exception as e:
            latency_ms = (time.time() - start_time) * 1000
            logger.error(f"[EnhancedPipeline] Retrieval error: {e}", exc_info=True)
            return self._create_empty_result(ctx, error=str(e), latency_ms=latency_ms)
    
    async def _hybrid_retrieve(
        self,
        tenant_id: str,
        query: str,
        embedding_result: Any,
        metadata_filter: dict[str, Any] | None = None,
        kb_ids: list[str] | None = None,
    ) -> list[RetrievalCandidate]:
        """混合检索：Dense + Keyword + RRF。"""
        client = await self._get_client()
        top_k = self._config.top_k
        expand_factor = self._config.hybrid.keyword_top_k_multiplier
        
        vector_task = self._dense_search(
            client=client,
            tenant_id=tenant_id,
            embedding=embedding_result.embedding_full,
            top_k=top_k * expand_factor,
            metadata_filter=metadata_filter,
            kb_ids=kb_ids,
        )
        
        keyword_task = self._keyword_search(
            client=client,
            tenant_id=tenant_id,
            query=query,
            top_k=top_k * expand_factor,
            metadata_filter=metadata_filter,
            kb_ids=kb_ids,
        ) if self._config.hybrid.enable_keyword else asyncio.sleep(0, result=[])
        
        vector_results, keyword_results = await asyncio.gather(
            vector_task, keyword_task, return_exceptions=True
        )
        
        if isinstance(vector_results, Exception):
            logger.warning(f"[EnhancedPipeline] Dense search failed: {vector_results}")
            vector_results = []
        
        if isinstance(keyword_results, Exception):
            logger.warning(f"[EnhancedPipeline] Keyword search failed: {keyword_results}")
            keyword_results = []
        
        combined = self._rrf_combiner.combine(
            vector_results=vector_results,
            bm25_results=keyword_results,
            vector_weight=self._config.hybrid.dense_weight,
            bm25_weight=self._config.hybrid.keyword_weight,
        )
        
        candidates = []
        for item in combined:
            candidates.append(RetrievalCandidate(
                id=item.get("id", ""),
                text=item.get("payload", {}).get("text", ""),
                score=item.get("score", 0.0),
                vector_score=item.get("vector_score", 0.0),
                keyword_score=item.get("bm25_score", 0.0),
                metadata=item.get("payload", {}),
            ))
        
        return candidates
    
    async def _dense_search(
        self,
        client: QdrantClient,
        tenant_id: str,
        embedding: list[float],
        top_k: int,
        metadata_filter: dict[str, Any] | None = None,
        kb_ids: list[str] | None = None,
    ) -> list[dict[str, Any]]:
        """Dense 向量检索。"""
        try:
            results = await client.search(
                tenant_id=tenant_id,
                query_vector=embedding,
                limit=top_k,
                vector_name="full",
                metadata_filter=metadata_filter,
                kb_ids=kb_ids,
            )
            return results
        except Exception as e:
            logger.error(f"[EnhancedPipeline] Dense search error: {e}")
            return []
    
    async def _keyword_search(
        self,
        client: QdrantClient,
        tenant_id: str,
        query: str,
        top_k: int,
        metadata_filter: dict[str, Any] | None = None,
        kb_ids: list[str] | None = None,
    ) -> list[dict[str, Any]]:
        """Keyword 关键词检索。"""
        try:
            qdrant = await client.get_client()
            collection_name = client.get_collection_name(tenant_id)
            
            query_terms = set(re.findall(r'\w+', query.lower()))
            
            results = await qdrant.scroll(
                collection_name=collection_name,
                limit=top_k * 3,
                with_payload=True,
            )
            
            scored_results = []
            for point in results[0]:
                text = point.payload.get("text", "").lower()
                text_terms = set(re.findall(r'\w+', text))
                overlap = len(query_terms & text_terms)
                
                if overlap > 0:
                    score = overlap / (len(query_terms) + len(text_terms) - overlap)
                    scored_results.append({
                        "id": str(point.id),
                        "score": score,
                        "payload": point.payload or {},
                    })
            
            scored_results.sort(key=lambda x: x["score"], reverse=True)
            return scored_results[:top_k]
            
        except Exception as e:
            logger.debug(f"[EnhancedPipeline] Keyword search failed: {e}")
            return []
    
    async def _rerank(
        self,
        candidates: list[RetrievalCandidate],
        query: str,
    ) -> list[RetrievalCandidate]:
        """可选重排。"""
        if not candidates:
            return candidates
        
        try:
            provider = await self._get_embedding_provider()
            query_embedding = await provider.embed_query(query)
            
            reranked = []
            for candidate in candidates:
                candidate_text = candidate.text[:500]
                if candidate_text:
                    candidate_embedding = await provider.embed(candidate_text)
                    similarity = self._cosine_similarity(
                        query_embedding.embedding_full,
                        candidate_embedding,
                    )
                    candidate.score = similarity
                
                if candidate.score >= self._reranker_config.min_score_threshold:
                    reranked.append(candidate)
            
            reranked.sort(key=lambda x: x.score, reverse=True)
            return reranked[:self._reranker_config.top_k_after_rerank]
            
        except Exception as e:
            logger.warning(f"[EnhancedPipeline] Rerank failed: {e}")
            return candidates
    
    def _cosine_similarity(self, vec1: list[float], vec2: list[float]) -> float:
        """计算余弦相似度。"""
        a = np.array(vec1)
        b = np.array(vec2)
        return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))
    
    async def health_check(self) -> bool:
        """健康检查。"""
        try:
            client = await self._get_client()
            qdrant = await client.get_client()
            await qdrant.get_collections()
            return True
        except Exception as e:
            logger.error(f"[EnhancedPipeline] Health check failed: {e}")
            return False


_enhanced_pipeline: EnhancedPipeline | None = None


async def get_enhanced_pipeline() -> EnhancedPipeline:
    """获取 EnhancedPipeline 单例。"""
    global _enhanced_pipeline
    if _enhanced_pipeline is None:
        _enhanced_pipeline = EnhancedPipeline()
    return _enhanced_pipeline
