"""Hybrid retriever — combines vector search (ChromaDB) + graph expansion (NetworkX) + community context."""

from __future__ import annotations

import logging
from typing import Optional

import networkx as nx

from res_sum.retrieval.vector_store import VectorStore
from res_sum.graph.knowledge_graph import KnowledgeGraph

logger = logging.getLogger("res_sum")


class HybridRetriever:
    """Retrieves relevant chunks using vector search, graph expansion, and community context.

    Supports four retrieval modes:
    - local: ChromaDB vector search only (best for specific factual queries)
    - graph: Graph traversal + chunk lookup (best for relational queries)
    - global: Community summaries + vector search (best for thematic synthesis)
    - hybrid: All three combined with re-ranking (default, best for general use)
    """

    def __init__(
        self,
        vector_store: VectorStore,
        knowledge_graph: Optional[KnowledgeGraph] = None,
        communities: Optional[dict[str, int]] = None,
        community_summaries: Optional[dict[int, str]] = None,
        alpha: float = 0.4,
        beta: float = 0.3,
        gamma: float = 0.3,
    ) -> None:
        """Initialize the hybrid retriever.

        Args:
            vector_store: ChromaDB vector store with ingested chunks.
            knowledge_graph: NetworkX knowledge graph (optional for local mode).
            communities: Entity → community ID mapping.
            community_summaries: Community ID → summary text.
            alpha: Weight for vector similarity score.
            beta: Weight for graph centrality score.
            gamma: Weight for community relevance score.
        """
        self.vector_store = vector_store
        self.kg = knowledge_graph
        self.communities = communities or {}
        self.community_summaries = community_summaries or {}
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma

    def retrieve(
        self,
        query: str,
        mode: str = "hybrid",
        n_results: int = 10,
        paper_filter: Optional[str] = None,
    ) -> list[dict]:
        """Retrieve relevant chunks based on the query and mode.

        Args:
            query: The search query.
            mode: Retrieval mode — "local", "graph", "global", or "hybrid".
            n_results: Number of results to return.
            paper_filter: If set, only retrieve from this paper.

        Returns:
            List of dicts with keys: chunk_id, text, score, source.
        """
        mode = mode.lower().strip()

        if mode == "local":
            return self._retrieve_local(query, n_results, paper_filter)
        elif mode == "graph":
            return self._retrieve_graph(query, n_results, paper_filter)
        elif mode == "global":
            return self._retrieve_global(query, n_results, paper_filter)
        elif mode == "hybrid":
            return self._retrieve_hybrid(query, n_results, paper_filter)
        else:
            raise ValueError(f"Unknown retrieval mode: {mode}. Use: local, graph, global, hybrid")

    # ------------------------------------------------------------------
    # Local: vector search only
    # ------------------------------------------------------------------

    def _retrieve_local(
        self, query: str, n_results: int, paper_filter: Optional[str]
    ) -> list[dict]:
        """Pure vector similarity search via ChromaDB."""
        where = {"paper_filename": paper_filter} if paper_filter else None
        results = self.vector_store.search(query, n_results=n_results, where=where)

        chunks = []
        for i, (doc, doc_id, distance) in enumerate(zip(
            results.get("documents", [[]])[0],
            results.get("ids", [[]])[0],
            results.get("distances", [[]])[0],
        )):
            # ChromaDB distances: lower = more similar. Convert to score.
            score = 1.0 / (1.0 + distance)
            chunks.append({
                "chunk_id": doc_id,
                "text": doc,
                "score": score,
                "source": "vector",
            })

        return chunks

    # ------------------------------------------------------------------
    # Graph: graph traversal + chunk lookup
    # ------------------------------------------------------------------

    def _retrieve_graph(
        self, query: str, n_results: int, paper_filter: Optional[str]
    ) -> list[dict]:
        """Graph-based retrieval: find entities, traverse neighbors, get their chunks."""
        if not self.kg or self.kg.num_nodes() == 0:
            logger.warning("No knowledge graph available, falling back to local retrieval")
            return self._retrieve_local(query, n_results, paper_filter)

        # Step 1: Vector search to find seed chunks
        seed_chunks = self._retrieve_local(query, n_results=5, paper_filter=paper_filter)
        seed_chunk_ids = {c["chunk_id"] for c in seed_chunks}

        # Step 2: Find entities connected to those chunks
        related_entities = set()
        for node, data in self.kg.graph.nodes(data=True):
            chunk_ids = set(data.get("source_chunk_ids", "").split("|"))
            if chunk_ids & seed_chunk_ids:
                related_entities.add(node)

        # Step 3: Expand to neighboring entities
        expanded_entities = set(related_entities)
        for entity in related_entities:
            expanded_entities.update(self.kg.neighbors(entity))

        # Step 4: Get chunk IDs from expanded entities
        expanded_chunk_ids = set()
        for entity in expanded_entities:
            expanded_chunk_ids.update(self.kg.get_source_chunk_ids_for_entity(entity))

        # Step 5: Fetch those chunks from ChromaDB
        new_chunk_ids = list(expanded_chunk_ids - seed_chunk_ids)
        graph_chunks = []
        if new_chunk_ids:
            try:
                fetched = self.vector_store.get_chunks_by_ids(new_chunk_ids[:n_results])
                for doc_id, doc in zip(
                    fetched.get("ids", []),
                    fetched.get("documents", []),
                ):
                    if doc:
                        graph_chunks.append({
                            "chunk_id": doc_id,
                            "text": doc,
                            "score": 0.5,  # Default score for graph-retrieved
                            "source": "graph",
                        })
            except Exception as e:
                logger.debug(f"Failed to fetch graph-expanded chunks: {e}")

        # Combine seed + graph chunks, deduplicate
        all_chunks = {c["chunk_id"]: c for c in seed_chunks}
        for c in graph_chunks:
            if c["chunk_id"] not in all_chunks:
                all_chunks[c["chunk_id"]] = c

        return sorted(all_chunks.values(), key=lambda x: x["score"], reverse=True)[:n_results]

    # ------------------------------------------------------------------
    # Global: community summaries + vector search
    # ------------------------------------------------------------------

    def _retrieve_global(
        self, query: str, n_results: int, paper_filter: Optional[str]
    ) -> list[dict]:
        """Community-based retrieval: find relevant communities, use their summaries as context."""
        chunks = self._retrieve_local(query, n_results=n_results, paper_filter=paper_filter)

        if not self.community_summaries:
            return chunks

        # Find which communities are represented in the retrieved chunks
        relevant_communities = set()
        if self.kg and self.communities:
            chunk_ids = {c["chunk_id"] for c in chunks}
            for node, data in self.kg.graph.nodes(data=True):
                node_chunks = set(data.get("source_chunk_ids", "").split("|"))
                if node_chunks & chunk_ids and node in self.communities:
                    relevant_communities.add(self.communities[node])

        # Add community summaries as additional context
        for comm_id in relevant_communities:
            if comm_id in self.community_summaries:
                chunks.append({
                    "chunk_id": f"community_{comm_id}",
                    "text": self.community_summaries[comm_id],
                    "score": 0.6,
                    "source": "community",
                })

        return chunks[:n_results + len(relevant_communities)]

    # ------------------------------------------------------------------
    # Hybrid: all three + re-ranking
    # ------------------------------------------------------------------

    def _retrieve_hybrid(
        self, query: str, n_results: int, paper_filter: Optional[str]
    ) -> list[dict]:
        """Combined retrieval with re-ranking."""
        # Gather candidates from all sources
        local_chunks = self._retrieve_local(query, n_results=n_results, paper_filter=paper_filter)
        graph_chunks = self._retrieve_graph(query, n_results=n_results, paper_filter=paper_filter)
        global_chunks = self._retrieve_global(query, n_results=n_results, paper_filter=paper_filter)

        # Deduplicate by chunk_id
        all_chunks: dict[str, dict] = {}
        for chunk in local_chunks + graph_chunks + global_chunks:
            cid = chunk["chunk_id"]
            if cid not in all_chunks or chunk["score"] > all_chunks[cid]["score"]:
                all_chunks[cid] = chunk

        # Re-rank using combined score
        if self.kg and self.kg.num_nodes() > 0:
            centrality = nx.degree_centrality(self.kg.graph)
            max_centrality = max(centrality.values()) if centrality else 1.0
        else:
            centrality = {}
            max_centrality = 1.0

        for chunk in all_chunks.values():
            vector_score = chunk["score"]

            # Graph centrality: average centrality of entities linked to this chunk
            graph_score = 0.0
            if self.kg and centrality:
                linked_entities = []
                for node, data in self.kg.graph.nodes(data=True):
                    if chunk["chunk_id"] in data.get("source_chunk_ids", ""):
                        linked_entities.append(node)
                if linked_entities:
                    graph_score = sum(
                        centrality.get(e, 0) for e in linked_entities
                    ) / (len(linked_entities) * max_centrality)

            # Community relevance: is this chunk from a well-represented community?
            community_score = 0.3 if chunk["source"] == "community" else 0.0

            chunk["score"] = (
                self.alpha * vector_score +
                self.beta * graph_score +
                self.gamma * community_score
            )

        ranked = sorted(all_chunks.values(), key=lambda x: x["score"], reverse=True)
        return ranked[:n_results]
