"""
MultiHopRetriever: Multi-hop reasoning and path ranking for Helix.

Implements BFS-based path exploration and scoring for complex queries
requiring multiple reasoning steps.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, NamedTuple
from collections import deque


class RetrievalPath(NamedTuple):
    """A path through the knowledge graph."""
    
    nodes: list[str]  # Entity IDs in order
    edges: list[dict]  # Edge properties between nodes
    total_score: float  # Combined path score
    hop_count: int  # Number of hops


@dataclass
class MultiHopRetriever:
    """
    Multi-hop retriever for complex reasoning queries.
    
    Uses BFS-based exploration with path scoring to find
    relevant reasoning chains in the knowledge graph.
    
    Key features:
    - BFS exploration from seed entities
    - Path scoring based on edge weights and relevance
    - Subgraph pruning for efficiency
    - Maximum hop limit to prevent explosion
    """
    
    # Graphiti instance for graph traversal
    graphiti: Any = None
    
    # Maximum number of hops to explore
    max_hops: int = 3
    
    # Maximum paths to return
    max_paths: int = 10
    
    # Minimum score threshold for paths
    min_score_threshold: float = 0.3
    
    # Edge weight decay per hop
    hop_decay: float = 0.8
    
    async def find_paths(
        self,
        query: str,
        seed_entities: list[str] | None = None,
        target_entities: list[str] | None = None,
    ) -> list[RetrievalPath]:
        """
        Find relevant paths in the knowledge graph.
        
        Args:
            query: The query string for relevance scoring.
            seed_entities: Starting entities (extracted from query if None).
            target_entities: Optional target entities to reach.
            
        Returns:
            List of scored paths through the KG.
        """
        if seed_entities is None:
            seed_entities = await self._extract_seed_entities(query)
        
        if not seed_entities:
            return []
        
        all_paths = []
        
        for seed in seed_entities:
            paths = await self._bfs_explore(
                seed,
                query=query,
                target_entities=target_entities,
            )
            all_paths.extend(paths)
        
        # Sort by score and limit
        all_paths.sort(key=lambda p: p.total_score, reverse=True)
        return all_paths[:self.max_paths]
    
    async def _extract_seed_entities(self, query: str) -> list[str]:
        """Extract seed entities from query using KG search."""
        if self.graphiti is None:
            return []
        
        try:
            results = await self.graphiti.search(
                query=query,
                num_results=5,
            )
            
            # Extract unique entity names from results
            entities = set()
            for result in results:
                # Graphiti returns edges with source/target nodes
                if hasattr(result, "source_node"):
                    entities.add(result.source_node.name)
                if hasattr(result, "target_node"):
                    entities.add(result.target_node.name)
            
            return list(entities)[:3]  # Limit seed entities
        except Exception:
            return []
    
    async def _bfs_explore(
        self,
        start_entity: str,
        query: str,
        target_entities: list[str] | None = None,
    ) -> list[RetrievalPath]:
        """
        BFS exploration from a starting entity.
        
        Args:
            start_entity: Starting entity name.
            query: Query for relevance scoring.
            target_entities: Optional targets to reach.
            
        Returns:
            List of discovered paths.
        """
        if self.graphiti is None:
            return []
        
        paths = []
        visited = {start_entity}
        
        # Queue: (current_path_nodes, current_path_edges, current_score)
        queue = deque([
            ([start_entity], [], 1.0)
        ])
        
        while queue:
            current_nodes, current_edges, current_score = queue.popleft()
            current_entity = current_nodes[-1]
            hop_count = len(current_edges)
            
            # Check hop limit
            if hop_count >= self.max_hops:
                # Save path if it has meaningful content
                if len(current_nodes) > 1:
                    paths.append(RetrievalPath(
                        nodes=current_nodes,
                        edges=current_edges,
                        total_score=current_score,
                        hop_count=hop_count,
                    ))
                continue
            
            # Get neighbors
            neighbors = await self._get_neighbors(current_entity)
            
            for neighbor, edge_data in neighbors:
                if neighbor in visited:
                    continue
                
                visited.add(neighbor)
                
                # Calculate edge score
                edge_score = await self._score_edge(
                    edge_data, query, hop_count
                )
                
                if edge_score < self.min_score_threshold:
                    continue
                
                # Create new path
                new_nodes = current_nodes + [neighbor]
                new_edges = current_edges + [edge_data]
                new_score = current_score * edge_score * self.hop_decay
                
                # Check if reached target
                if target_entities and neighbor in target_entities:
                    paths.append(RetrievalPath(
                        nodes=new_nodes,
                        edges=new_edges,
                        total_score=new_score,
                        hop_count=hop_count + 1,
                    ))
                else:
                    queue.append((new_nodes, new_edges, new_score))
        
        return paths
    
    async def _get_neighbors(self, entity: str) -> list[tuple[str, dict]]:
        """
        Get neighboring entities and edges.
        
        Returns:
            List of (neighbor_name, edge_data) tuples.
        """
        if self.graphiti is None:
            return []
        
        try:
            # Query Neo4j for neighbors
            query = """
            MATCH (n:Entity {name: $name})-[r]-(m:Entity)
            RETURN m.name as neighbor, 
                   type(r) as relation_type,
                   r.fact as fact,
                   r.weight as weight
            LIMIT 20
            """
            results = await self.graphiti.driver.execute_query(
                query,
                name=entity,
            )
            
            neighbors = []
            for result in results:
                edge_data = {
                    "relation_type": result.get("relation_type", "RELATED_TO"),
                    "fact": result.get("fact", ""),
                    "weight": result.get("weight", 1.0),
                }
                neighbors.append((result["neighbor"], edge_data))
            
            return neighbors
        except Exception:
            return []
    
    async def _score_edge(
        self,
        edge_data: dict,
        query: str,
        hop_count: int,
    ) -> float:
        """
        Score an edge's relevance to the query.
        
        Args:
            edge_data: Edge properties.
            query: Query string.
            hop_count: Current hop count (for decay).
            
        Returns:
            Relevance score between 0 and 1.
        """
        # Base weight from edge
        base_weight = edge_data.get("weight", 1.0)
        
        # Text relevance (simple keyword overlap)
        fact = edge_data.get("fact", "").lower()
        query_terms = set(query.lower().split())
        fact_terms = set(fact.split())
        
        overlap = len(query_terms & fact_terms)
        text_score = min(overlap / max(len(query_terms), 1), 1.0)
        
        # Combine scores
        combined = 0.5 * base_weight + 0.5 * text_score
        
        return min(combined, 1.0)
    
    def format_paths_as_context(self, paths: list[RetrievalPath]) -> str:
        """
        Format paths as context for LLM.
        
        Args:
            paths: List of retrieval paths.
            
        Returns:
            Formatted string for LLM context.
        """
        if not paths:
            return ""
        
        parts = ["**Reasoning Paths from Knowledge Graph:**\n"]
        
        for i, path in enumerate(paths, 1):
            parts.append(f"\n**Path {i}** (score: {path.total_score:.2f}, hops: {path.hop_count}):")
            
            for j, (node, edge) in enumerate(zip(path.nodes[:-1], path.edges)):
                next_node = path.nodes[j + 1]
                relation = edge.get("relation_type", "→")
                fact = edge.get("fact", "")
                
                parts.append(f"  {node} --[{relation}]--> {next_node}")
                if fact:
                    parts.append(f"    Fact: {fact[:100]}...")
        
        return "\n".join(parts)
