"""Context selector using graph-based analysis."""

import logging
import time
from collections import defaultdict, deque
from typing import Any, Dict, List, Optional, Set

from .interfaces import (
    EdgeType,
    GraphCutResult,
    GraphEdge,
    GraphNode,
    IContextSelector,
)

logger = logging.getLogger(__name__)

# Check if TreeSitter Chunker graph cut is available
GRAPH_CUT_AVAILABLE = False
try:
    from chunker.graph.cut import graph_cut

    GRAPH_CUT_AVAILABLE = True
    logger.info("TreeSitter Chunker graph_cut available")
except ImportError:
    logger.warning("TreeSitter Chunker graph_cut not available, using fallback")


class ContextSelector(IContextSelector):
    """Selects optimal context using graph analysis."""

    def __init__(self, nodes: List[GraphNode], edges: List[GraphEdge]):
        """
        Initialize the context selector.

        Args:
            nodes: Graph nodes
            edges: Graph edges
        """
        self.nodes = nodes
        self.edges = edges
        self._node_map: Dict[str, GraphNode] = {n.id: n for n in nodes}

        logger.debug(f"ContextSelector initialized: {len(nodes)} nodes, {len(edges)} edges")

    def select_context(
        self,
        seeds: List[str],
        radius: int = 2,
        budget: int = 200,
        weights: Optional[Dict[str, float]] = None,
    ) -> GraphCutResult:
        """
        Select optimal context around seed nodes.

        Uses TreeSitter Chunker's graph_cut if available, otherwise falls back to BFS.

        Args:
            seeds: Seed node IDs
            radius: Maximum distance from seeds
            budget: Maximum number of nodes to select
            weights: Scoring weights for graph cut

        Returns:
            GraphCutResult with selected nodes and edges
        """
        start_time = time.time()

        if GRAPH_CUT_AVAILABLE:
            result = self._select_with_graph_cut(seeds, radius, budget, weights)
        else:
            result = self._select_with_bfs(seeds, radius, budget)

        execution_time = (time.time() - start_time) * 1000  # Convert to ms
        result.execution_time_ms = execution_time

        logger.info(
            f"Selected {len(result.selected_nodes)} nodes in {execution_time:.2f}ms "
            f"(seeds={len(seeds)}, radius={radius}, budget={budget})"
        )

        return result

    def _select_with_graph_cut(
        self,
        seeds: List[str],
        radius: int,
        budget: int,
        weights: Optional[Dict[str, float]],
    ) -> GraphCutResult:
        """
        Select context using TreeSitter Chunker's graph_cut.

        Args:
            seeds: Seed node IDs
            radius: Maximum distance from seeds
            budget: Maximum number of nodes to select
            weights: Scoring weights

        Returns:
            GraphCutResult
        """
        # Convert nodes and edges to chunker format
        raw_nodes = []
        for node in self.nodes:
            raw_nodes.append(
                {
                    "id": node.id,
                    "file": node.file_path,
                    "lang": node.language,
                    "symbol": node.symbol,
                    "kind": node.kind,
                    "attrs": node.attrs,
                }
            )

        raw_edges = []
        for edge in self.edges:
            raw_edges.append(
                {
                    "src": edge.source_id,
                    "dst": edge.target_id,
                    "type": edge.edge_type.value,
                    "weight": edge.weight,
                }
            )

        # Call graph_cut
        selected_ids, induced_edges_raw = graph_cut(
            seeds=seeds,
            nodes=raw_nodes,
            edges=raw_edges,
            radius=radius,
            budget=budget,
            weights=weights or {},
        )

        # Convert results back to our format
        selected_nodes = [self._node_map[nid] for nid in selected_ids if nid in self._node_map]

        induced_edges = []
        for raw_edge in induced_edges_raw:
            edge_type_str = raw_edge.get("type", "REFERENCES").upper()
            try:
                edge_type = EdgeType[edge_type_str]
            except KeyError:
                edge_type = EdgeType.REFERENCES

            induced_edges.append(
                GraphEdge(
                    source_id=raw_edge.get("src", ""),
                    target_id=raw_edge.get("dst", ""),
                    edge_type=edge_type,
                    weight=raw_edge.get("weight", 1.0),
                )
            )

        return GraphCutResult(
            selected_nodes=selected_nodes,
            induced_edges=induced_edges,
            seed_nodes=seeds,
            radius=radius,
            budget=budget,
            total_candidates=len(raw_nodes),
            execution_time_ms=0.0,  # Will be set by caller
        )

    def _select_with_bfs(self, seeds: List[str], radius: int, budget: int) -> GraphCutResult:
        """
        Fallback BFS-based context selection.

        Args:
            seeds: Seed node IDs
            radius: Maximum distance from seeds
            budget: Maximum number of nodes to select

        Returns:
            GraphCutResult
        """
        # Build adjacency lists
        out_adj: Dict[str, Set[str]] = defaultdict(set)
        out_degree: Dict[str, int] = defaultdict(int)

        for edge in self.edges:
            out_adj[edge.source_id].add(edge.target_id)
            out_degree[edge.source_id] += 1

        # BFS from seeds to collect candidates within radius
        distance: Dict[str, int] = {}
        queue: deque = deque()

        for seed in seeds:
            if seed in self._node_map:
                distance[seed] = 0
                queue.append(seed)

        visited: Set[str] = set(distance.keys())

        while queue:
            current = queue.popleft()
            current_dist = distance[current]

            if current_dist >= radius:
                continue

            for neighbor in out_adj.get(current, set()):
                if neighbor not in visited:
                    visited.add(neighbor)
                    distance[neighbor] = current_dist + 1
                    queue.append(neighbor)

        # Score candidates by distance and degree
        def score(node_id: str) -> float:
            dist = distance.get(node_id, 1000000)
            if dist <= 0:
                dist = 1
            degree = out_degree.get(node_id, 0)
            # Simple scoring: closer nodes and higher degree are better
            return (1.0 / float(dist)) + (degree * 0.1)

        candidates = [nid for nid in visited if nid in self._node_map]
        candidates.sort(key=score, reverse=True)

        # Select top budget nodes
        selected_ids = candidates[:budget]
        selected_set = set(selected_ids)

        selected_nodes = [self._node_map[nid] for nid in selected_ids]

        # Induced edges
        induced_edges = [
            edge
            for edge in self.edges
            if edge.source_id in selected_set and edge.target_id in selected_set
        ]

        return GraphCutResult(
            selected_nodes=selected_nodes,
            induced_edges=induced_edges,
            seed_nodes=seeds,
            radius=radius,
            budget=budget,
            total_candidates=len(candidates),
            execution_time_ms=0.0,  # Will be set by caller
        )

    def expand_search_results(
        self,
        search_results: List[Dict[str, Any]],
        expansion_radius: int = 1,
        max_context_nodes: int = 50,
    ) -> List[GraphNode]:
        """
        Expand search results with graph context.

        Args:
            search_results: Search results to expand (must have 'file' key)
            expansion_radius: How far to expand
            max_context_nodes: Maximum context nodes to add

        Returns:
            List of context nodes
        """
        # Extract file paths from search results
        result_files = set()
        for result in search_results:
            file_path = result.get("file") or result.get("file_path")
            if file_path:
                result_files.add(file_path)

        if not result_files:
            logger.warning("No file paths in search results")
            return []

        # Find nodes matching these files
        seed_nodes = []
        for node in self.nodes:
            if node.file_path in result_files:
                seed_nodes.append(node.id)

        if not seed_nodes:
            logger.warning(f"No graph nodes found for {len(result_files)} result files")
            return []

        # Use graph cut to select context
        result = self.select_context(
            seeds=seed_nodes,
            radius=expansion_radius,
            budget=max_context_nodes,
        )

        logger.debug(
            f"Expanded {len(search_results)} search results to {len(result.selected_nodes)} context nodes"
        )

        return result.selected_nodes
