"""
Hybrid Retriever - Intelligent context retrieval combining multiple strategies

Combines:
1. BM25 search for keyword matching
2. Code graph traversal for related code
3. Hierarchical navigation for context expansion
4. Score fusion and re-ranking for optimal results

This replaces simple keyword matching with a multi-signal retrieval approach
that understands code structure and relationships.
"""

from dataclasses import dataclass
from typing import Any, Optional

from groknroll.oracle.bm25_index import BM25Index
from groknroll.oracle.code_graph import CodeGraphBuilder
from groknroll.storage.database import Database


@dataclass
class RetrievalResult:
    """Result from hybrid retrieval"""

    summary_id: int
    score: float
    scope_type: str
    scope_path: str
    scope_name: Optional[str]
    summary: Optional[str]
    keywords: Optional[list[str]]
    source: str  # 'bm25', 'graph', 'hierarchy', 'combined'
    content: Optional[str] = None  # File content if loaded


class HybridRetriever:
    """
    Intelligent context retrieval combining multiple strategies.

    Retrieval pipeline:
    1. BM25 search for initial keyword matches
    2. Graph expansion for related code (imports, inheritance, calls)
    3. Hierarchical expansion (parent/child summaries)
    4. Score fusion with configurable weights
    5. Re-ranking based on relevance signals
    """

    # Score fusion weights
    BM25_WEIGHT = 0.6
    GRAPH_WEIGHT = 0.3
    HIERARCHY_WEIGHT = 0.1

    # Retrieval limits
    MAX_BM25_RESULTS = 30
    MAX_GRAPH_DEPTH = 2
    MAX_GRAPH_EXPANSION = 20

    def __init__(
        self,
        database: Database,
        project_id: int,
        bm25_index: BM25Index,
        code_graph: CodeGraphBuilder,
        verbose: bool = False,
    ):
        """
        Initialize hybrid retriever.

        Args:
            database: Database instance
            project_id: Project ID
            bm25_index: BM25 index for keyword search
            code_graph: Code graph for relationship traversal
            verbose: Whether to print progress
        """
        self.database = database
        self.project_id = project_id
        self.bm25_index = bm25_index
        self.code_graph = code_graph
        self.verbose = verbose

        # Cache for file contents
        self._content_cache: dict[str, str] = {}

    def retrieve(
        self,
        query: str,
        max_results: int = 10,
        include_content: bool = True,
        scope_types: Optional[list[str]] = None,
        expand_graph: bool = True,
        expand_hierarchy: bool = True,
    ) -> list[RetrievalResult]:
        """
        Retrieve relevant code using hybrid approach.

        Args:
            query: Search query
            max_results: Maximum results to return
            include_content: Whether to load file contents
            scope_types: Filter by scope types (optional)
            expand_graph: Whether to expand via code graph
            expand_hierarchy: Whether to expand via hierarchy

        Returns:
            List of RetrievalResult sorted by relevance
        """
        if self.verbose:
            print(f"🔎 Retrieving context for: {query[:50]}...")

        # Step 1: BM25 search for keyword matches
        bm25_results = self._bm25_search(query, scope_types)

        if self.verbose:
            print(f"   BM25: {len(bm25_results)} results")

        # Step 2: Graph expansion for related code
        graph_results = []
        if expand_graph and bm25_results:
            graph_results = self._graph_expansion(bm25_results)
            if self.verbose:
                print(f"   Graph expansion: {len(graph_results)} related")

        # Step 3: Hierarchical expansion
        hierarchy_results = []
        if expand_hierarchy and bm25_results:
            hierarchy_results = self._hierarchy_expansion(bm25_results)
            if self.verbose:
                print(f"   Hierarchy expansion: {len(hierarchy_results)} related")

        # Step 4: Score fusion
        fused_results = self._fuse_scores(
            bm25_results, graph_results, hierarchy_results
        )

        if self.verbose:
            print(f"   After fusion: {len(fused_results)} unique results")

        # Step 5: Re-rank and limit
        ranked_results = self._rerank(fused_results, query)[:max_results]

        # Step 6: Load content if requested
        if include_content:
            self._load_contents(ranked_results)

        if self.verbose:
            print(f"✅ Retrieved {len(ranked_results)} results")

        return ranked_results

    def retrieve_for_file(
        self,
        file_path: str,
        max_results: int = 10,
    ) -> list[RetrievalResult]:
        """
        Retrieve related code for a specific file.

        Args:
            file_path: Path to the file
            max_results: Maximum results to return

        Returns:
            List of RetrievalResult for related files
        """
        # Get summary for this file
        summary = self.database.get_summary(
            self.project_id, "module", file_path
        )

        if not summary:
            return []

        # Get related via graph
        related_nodes = self.code_graph.get_related(
            summary.id,
            max_depth=self.MAX_GRAPH_DEPTH,
        )

        results = []
        for node in related_nodes[:max_results]:
            node_summary = self.database.get_summary_by_id(node.summary_id)
            if node_summary:
                # Score decreases with depth
                score = 1.0 / (1 + node.depth)
                results.append(RetrievalResult(
                    summary_id=node.summary_id,
                    score=score,
                    scope_type=node.scope_type,
                    scope_path=node.scope_path,
                    scope_name=node.scope_name,
                    summary=node_summary.summary,
                    keywords=node_summary.keywords,
                    source="graph",
                ))

        return sorted(results, key=lambda r: -r.score)

    def _bm25_search(
        self,
        query: str,
        scope_types: Optional[list[str]] = None,
    ) -> list[RetrievalResult]:
        """Perform BM25 search."""
        bm25_results = self.bm25_index.search(
            query,
            limit=self.MAX_BM25_RESULTS,
            scope_types=scope_types,
        )

        return [
            RetrievalResult(
                summary_id=r.summary_id,
                score=r.score,
                scope_type=r.scope_type,
                scope_path=r.scope_path,
                scope_name=r.scope_name,
                summary=r.summary,
                keywords=r.keywords,
                source="bm25",
            )
            for r in bm25_results
        ]

    def _graph_expansion(
        self,
        seed_results: list[RetrievalResult],
    ) -> list[RetrievalResult]:
        """Expand results using code graph."""
        expanded = []
        seen_ids = {r.summary_id for r in seed_results}

        for seed in seed_results[:10]:  # Limit seeds for performance
            related = self.code_graph.get_related(
                seed.summary_id,
                max_depth=self.MAX_GRAPH_DEPTH,
                relation_types=["imports", "inherits", "calls"],
            )

            for node in related:
                if node.summary_id not in seen_ids:
                    seen_ids.add(node.summary_id)

                    summary = self.database.get_summary_by_id(node.summary_id)
                    if summary:
                        # Score based on seed score and depth
                        score = seed.score * (0.5 ** node.depth)
                        expanded.append(RetrievalResult(
                            summary_id=node.summary_id,
                            score=score,
                            scope_type=node.scope_type,
                            scope_path=node.scope_path,
                            scope_name=node.scope_name,
                            summary=summary.summary,
                            keywords=summary.keywords,
                            source="graph",
                        ))

                    if len(expanded) >= self.MAX_GRAPH_EXPANSION:
                        break

            if len(expanded) >= self.MAX_GRAPH_EXPANSION:
                break

        return expanded

    def _hierarchy_expansion(
        self,
        seed_results: list[RetrievalResult],
    ) -> list[RetrievalResult]:
        """Expand results using summary hierarchy."""
        expanded = []
        seen_ids = {r.summary_id for r in seed_results}

        for seed in seed_results[:10]:
            summary = self.database.get_summary_by_id(seed.summary_id)
            if not summary:
                continue

            # Get parent
            if summary.parent_id and summary.parent_id not in seen_ids:
                parent = self.database.get_summary_by_id(summary.parent_id)
                if parent:
                    seen_ids.add(parent.id)
                    expanded.append(RetrievalResult(
                        summary_id=parent.id,
                        score=seed.score * 0.7,
                        scope_type=parent.scope_type,
                        scope_path=parent.scope_path,
                        scope_name=parent.scope_name,
                        summary=parent.summary,
                        keywords=parent.keywords,
                        source="hierarchy",
                    ))

            # Get children
            children = self.database.get_child_summaries(summary.id)
            for child in children[:5]:
                if child.id not in seen_ids:
                    seen_ids.add(child.id)
                    expanded.append(RetrievalResult(
                        summary_id=child.id,
                        score=seed.score * 0.5,
                        scope_type=child.scope_type,
                        scope_path=child.scope_path,
                        scope_name=child.scope_name,
                        summary=child.summary,
                        keywords=child.keywords,
                        source="hierarchy",
                    ))

        return expanded

    def _fuse_scores(
        self,
        bm25_results: list[RetrievalResult],
        graph_results: list[RetrievalResult],
        hierarchy_results: list[RetrievalResult],
    ) -> list[RetrievalResult]:
        """Fuse scores from multiple sources."""
        # Normalize scores within each source
        bm25_results = self._normalize_scores(bm25_results)
        graph_results = self._normalize_scores(graph_results)
        hierarchy_results = self._normalize_scores(hierarchy_results)

        # Combine into single dict by summary_id
        combined: dict[int, RetrievalResult] = {}

        for result in bm25_results:
            combined[result.summary_id] = RetrievalResult(
                summary_id=result.summary_id,
                score=result.score * self.BM25_WEIGHT,
                scope_type=result.scope_type,
                scope_path=result.scope_path,
                scope_name=result.scope_name,
                summary=result.summary,
                keywords=result.keywords,
                source="combined",
            )

        for result in graph_results:
            if result.summary_id in combined:
                combined[result.summary_id].score += result.score * self.GRAPH_WEIGHT
            else:
                combined[result.summary_id] = RetrievalResult(
                    summary_id=result.summary_id,
                    score=result.score * self.GRAPH_WEIGHT,
                    scope_type=result.scope_type,
                    scope_path=result.scope_path,
                    scope_name=result.scope_name,
                    summary=result.summary,
                    keywords=result.keywords,
                    source="combined",
                )

        for result in hierarchy_results:
            if result.summary_id in combined:
                combined[result.summary_id].score += result.score * self.HIERARCHY_WEIGHT
            else:
                combined[result.summary_id] = RetrievalResult(
                    summary_id=result.summary_id,
                    score=result.score * self.HIERARCHY_WEIGHT,
                    scope_type=result.scope_type,
                    scope_path=result.scope_path,
                    scope_name=result.scope_name,
                    summary=result.summary,
                    keywords=result.keywords,
                    source="combined",
                )

        return list(combined.values())

    def _normalize_scores(
        self,
        results: list[RetrievalResult],
    ) -> list[RetrievalResult]:
        """Normalize scores to 0-1 range."""
        if not results:
            return results

        max_score = max(r.score for r in results)
        if max_score == 0:
            return results

        for result in results:
            result.score = result.score / max_score

        return results

    def _rerank(
        self,
        results: list[RetrievalResult],
        query: str,
    ) -> list[RetrievalResult]:
        """Re-rank results based on additional signals."""
        # Apply boosting based on scope type
        scope_boosts = {
            "module": 1.2,   # Prefer file-level results
            "class": 1.1,   # Then classes
            "function": 1.0, # Then functions
            "package": 0.8,  # Packages less specific
        }

        for result in results:
            boost = scope_boosts.get(result.scope_type, 1.0)
            result.score *= boost

            # Boost if query terms appear in scope name
            if result.scope_name:
                query_lower = query.lower()
                name_lower = result.scope_name.lower()
                if name_lower in query_lower or query_lower in name_lower:
                    result.score *= 1.5

        # Sort by score descending
        return sorted(results, key=lambda r: -r.score)

    def _load_contents(self, results: list[RetrievalResult]) -> None:
        """Load file contents for results."""
        for result in results:
            if result.scope_type in ("module", "class", "function"):
                content = self._get_file_content(result.scope_path)
                result.content = content

    def _get_file_content(self, file_path: str) -> Optional[str]:
        """Get file content from cache or disk."""
        if file_path in self._content_cache:
            return self._content_cache[file_path]

        # Try to read from database summaries
        summary = self.database.get_summary(self.project_id, "module", file_path)
        if summary:
            # Content not stored in summary, would need to read from disk
            # This would require access to the original indexer
            pass

        return None

    def set_content_cache(self, files: dict[str, Any]) -> None:
        """
        Set content cache from indexed files.

        Args:
            files: Dict mapping relative_path to FileInfo from indexer
        """
        for rel_path, file_info in files.items():
            if file_info.content:
                self._content_cache[rel_path] = file_info.content

    def clear_cache(self) -> None:
        """Clear content cache."""
        self._content_cache.clear()


def create_hybrid_retriever(
    database: Database,
    project_id: int,
    verbose: bool = False,
) -> HybridRetriever:
    """
    Factory function to create a configured HybridRetriever.

    Args:
        database: Database instance
        project_id: Project ID
        verbose: Whether to print progress

    Returns:
        Configured HybridRetriever instance
    """
    bm25_index = BM25Index(database, project_id, verbose=verbose)
    code_graph = CodeGraphBuilder(database, project_id, verbose=verbose)

    return HybridRetriever(
        database=database,
        project_id=project_id,
        bm25_index=bm25_index,
        code_graph=code_graph,
        verbose=verbose,
    )
