"""
ChunkHive Retrieval Engine
===========================

Full pipeline:
Embedding → FAISS → Ranking → Graph Expansion → Context Building
"""

from typing import List, Any

from chunkhive.embeddings.embedding_engine import SemanticEmbeddingEngine
from chunkhive.retrieval.vector_store import VectorStore
from chunkhive.retrieval.ranking_engine import RankingEngine
from chunkhive.retrieval.context_engine import ContextEngine

from chunkhive.graph.relationship_graph import RelationshipGraph
from chunkhive.graph.dependency_expander import DependencyExpander

from chunkhive.retrieval.scoring.semantic import SemanticScorer
from chunkhive.retrieval.scoring.keyword import KeywordScorer
from chunkhive.retrieval.scoring.structure import StructuralScorer
from chunkhive.retrieval.scoring.dependency import DependencyScorer
from chunkhive.retrieval.scoring.metadata import MetadataScorer
from chunkhive.retrieval.scoring.aggregator import AggregatorScorer


class RetrievalEngine:

    def __init__(
        self,
        vector_store: VectorStore,
        model_name: str = "minilm",
        enable_expansion: bool = True,
    ):
        self.vector_store = vector_store
        self.embedding_engine = SemanticEmbeddingEngine(model_name=model_name)
        self.enable_expansion = enable_expansion

        # -------- Build Scorers --------
        self.semantic_scorer = SemanticScorer(
            embedding_engine=self.embedding_engine,
            weight=1.0
        )

        self.keyword_scorer = KeywordScorer(weight=0.3)
        self.structural_scorer = StructuralScorer(weight=0.2)
        self.dependency_scorer = DependencyScorer(weight=0.3)
        self.metadata_scorer = MetadataScorer()

        self.aggregator = AggregatorScorer(
            scorers=[
                self.semantic_scorer,
                self.keyword_scorer,
                self.structural_scorer,
                self.dependency_scorer,
                self.metadata_scorer,
            ]
        )

        self.ranking_engine = RankingEngine(
            aggregator=self.aggregator
        )

    # ============================================================
    # SEARCH PIPELINE
    # ============================================================

    def search(self, query: str, top_k: int = 5) -> List[Any]:

        # 1️⃣ Encode query
        query_embedding = self.embedding_engine.generate_embedding_for_query(query)

        # 2️⃣ FAISS retrieval (broad pool)
        vector_results = self.vector_store.search(
            query_vector=query_embedding,
            top_k=top_k * 4,
        )

        if not vector_results:
            return []

        # Extract chunks
        candidates = [chunk for chunk, _ in vector_results]

        # 3️⃣ Prepare semantic scorer once
        self.semantic_scorer.prepare_query(query)

        # 4️⃣ Rank candidates
        ranked_results = self.ranking_engine.rank(
            query=query,
            candidates=candidates,
            top_k=top_k,
            explain=False,
        )

        # 5️⃣ Graph expansion
        graph = RelationshipGraph()
        graph.build(candidates)

        expander = None
        if self.enable_expansion:
            expander = DependencyExpander(graph=graph)

        # 6️⃣ Build context ONCE (not inside loop!)
        context_engine = ContextEngine(
            ranking_engine=self.ranking_engine,
            dependency_expander=expander,
            max_context_chunks=8,
            max_total_chars=12000,
        )

        # Build context once - it already uses ranked_results internally
        context_chunks = context_engine.build_context(
            query=query,
            candidates=candidates,
            top_k=top_k,
            expand=self.enable_expansion,
        )

        # 7️⃣ Attach context to each ranked result
        final_results = []
        for result in ranked_results:
            final_results.append({
                "score": result.score,
                "confidence": result.confidence,  # Add confidence!
                "chunk": result.chunk,
                "context": context_chunks,  # Same context for all
            })

        return final_results

        # # 6️⃣ Context building
        # context_engine = ContextEngine(
        #     ranking_engine=self.ranking_engine,
        #     dependency_expander=expander,
        #     max_context_chunks=8,
        #     max_total_chars=12000,
        # )

        # final_results = []

        # for result in ranked_results:

        #     context_chunks = context_engine.build_context(
        #         query=query,
        #         candidates=candidates,
        #         top_k=top_k,
        #         expand=self.enable_expansion,
        #     )

        #     final_results.append({
        #         "score": result.score,
        #         "chunk": result.chunk,
        #         "context": context_chunks,
        #     })

        # return final_results
