"""
Tool Relationship Graph

Provides a rustworkx PyDiGraph tracking tool co-occurrence patterns
across sessions. Used by BehavioralRelationshipStrategy to generate
relationship text for the RIC embedding pipeline.

Nodes: One per unique tool_name (~50 tools)
Edges: Temporal succession within sessions (A -> B = A called before B)

Graph analysis (all via rustworkx):
    - k-core membership per tool
    - betweenness/closeness centrality
    - predecessor/successor frequency
    - cycle detection for recurring tool loops
"""

import logging
import time
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Set, Tuple

logger = logging.getLogger(__name__)

# Lazy rustworkx import (reuse pattern from graph_mixin.py)
_rustworkx = None


def _get_rustworkx():
    """Lazy load rustworkx."""
    global _rustworkx
    if _rustworkx is None:
        import rustworkx as rx

        _rustworkx = rx
    return _rustworkx


@dataclass
class ToolNodeData:
    """Attributes stored on each tool node."""

    tool_name: str
    service: str = ""
    first_seen: float = 0.0
    call_count: int = 0
    total_execution_ms: float = 0.0

    @property
    def avg_execution_ms(self) -> float:
        if self.call_count == 0:
            return 0.0
        return self.total_execution_ms / self.call_count


@dataclass
class ToolEdgeData:
    """Attributes stored on each co-occurrence edge."""

    co_occurrence_count: int = 0
    total_time_delta_ms: float = 0.0
    sessions_seen_in: Set[str] = field(default_factory=set)

    @property
    def avg_time_delta_ms(self) -> float:
        if self.co_occurrence_count == 0:
            return 0.0
        return self.total_time_delta_ms / self.co_occurrence_count


class ToolRelationshipGraph:
    """
    Rustworkx-based directed graph tracking tool co-occurrence.

    Usage:
        graph = ToolRelationshipGraph()

        # Record tool calls in a session
        graph.record_tool_call("list_spaces", "chat", session_id="s1")
        graph.record_tool_call("send_message", "chat", session_id="s1")

        # Get relationship text for embedding
        text = graph.get_relationship_text("send_message", user_email="u@x.com", session_id="s1")

        # Recompute expensive metrics periodically
        graph.recompute_metrics()
    """

    # How many recent tool calls per session to link as predecessors
    PREDECESSOR_WINDOW = 3

    def __init__(self):
        rx = _get_rustworkx()
        self._graph: Any = rx.PyDiGraph()  # rx.PyDiGraph

        # Bidirectional name <-> index mappings
        self._name_to_idx: Dict[str, int] = {}
        self._idx_to_name: Dict[int, str] = {}

        # Per-session recent tool history (for edge creation)
        self._session_history: Dict[str, List[Tuple[str, float]]] = defaultdict(list)

        # Cached graph metrics (recomputed periodically)
        self._betweenness: Dict[int, float] = {}
        self._closeness: Dict[int, float] = {}
        self._k_cores: Dict[int, int] = {}
        self._metrics_stale = True

    # =========================================================================
    # NODE / EDGE MANAGEMENT
    # =========================================================================

    def _ensure_node(self, tool_name: str, service: str = "") -> int:
        """Add a tool node if it doesn't exist. Returns node index."""
        if tool_name in self._name_to_idx:
            return self._name_to_idx[tool_name]

        data = ToolNodeData(
            tool_name=tool_name,
            service=service,
            first_seen=time.time(),
        )
        idx = self._graph.add_node(data)
        self._name_to_idx[tool_name] = idx
        self._idx_to_name[idx] = tool_name
        return idx

    def _ensure_edge(self, src_idx: int, dst_idx: int) -> ToolEdgeData:
        """Add or get the edge between two nodes."""
        if self._graph.has_edge(src_idx, dst_idx):
            edge_idx = self._graph.get_edge_data(src_idx, dst_idx)
            # edge_idx is actually the edge data object
            return edge_idx
        data = ToolEdgeData()
        self._graph.add_edge(src_idx, dst_idx, data)
        return data

    # =========================================================================
    # RECORDING
    # =========================================================================

    def record_tool_call(
        self,
        tool_name: str,
        service: str = "",
        session_id: str = "",
        execution_time_ms: float = 0.0,
        timestamp_ms: Optional[float] = None,
    ) -> None:
        """
        Record a tool call and update the co-occurrence graph.

        Creates edges from recent prior tools in the same session to this tool.

        Args:
            tool_name: Name of the tool called
            service: Service group (e.g., "chat", "drive")
            session_id: Session identifier for temporal grouping
            execution_time_ms: How long the tool took
            timestamp_ms: Call timestamp (defaults to now)
        """
        now = timestamp_ms or (time.time() * 1000)
        idx = self._ensure_node(tool_name, service)

        # Update node stats
        node_data: ToolNodeData = self._graph.get_node_data(idx)
        node_data.call_count += 1
        node_data.total_execution_ms += execution_time_ms
        if service and not node_data.service:
            node_data.service = service

        # Create edges from recent predecessors in this session
        if session_id:
            history = self._session_history[session_id]

            # Link from each recent predecessor to this tool
            window = history[-self.PREDECESSOR_WINDOW :]
            for prev_name, prev_ts in window:
                if prev_name == tool_name:
                    continue  # Skip self-loops
                prev_idx = self._name_to_idx.get(prev_name)
                if prev_idx is None:
                    continue
                edge_data = self._ensure_edge(prev_idx, idx)
                edge_data.co_occurrence_count += 1
                edge_data.total_time_delta_ms += now - prev_ts
                edge_data.sessions_seen_in.add(session_id)

            # Append to session history
            history.append((tool_name, now))

            # Trim old history to prevent unbounded growth
            if len(history) > 50:
                self._session_history[session_id] = history[-50:]

        self._metrics_stale = True

    # =========================================================================
    # GRAPH ANALYSIS
    # =========================================================================

    def recompute_metrics(self) -> None:
        """Recompute expensive graph metrics (centrality, k-cores).

        Call periodically (e.g., every N tool calls or on a timer),
        not on every record_tool_call.
        """
        rx = _get_rustworkx()
        graph = self._graph

        if graph.num_nodes() == 0:
            self._betweenness = {}
            self._closeness = {}
            self._k_cores = {}
            self._metrics_stale = False
            return

        try:
            self._betweenness = rx.digraph_betweenness_centrality(graph)
        except Exception:
            self._betweenness = {}

        try:
            self._closeness = rx.digraph_closeness_centrality(graph)
        except Exception:
            self._closeness = {}

        try:
            # k-core requires undirected graph
            undirected = graph.to_undirected()
            self._k_cores = rx.core_number(undirected)
        except Exception:
            self._k_cores = {}

        self._metrics_stale = False
        logger.debug(
            f"Recomputed graph metrics: {graph.num_nodes()} nodes, "
            f"{graph.num_edges()} edges"
        )

    def get_predecessors(self, tool_name: str, top_n: int = 3) -> List[Tuple[str, int]]:
        """Get top-N predecessors by co-occurrence count.

        Returns:
            List of (tool_name, count) tuples sorted by count descending.
        """
        idx = self._name_to_idx.get(tool_name)
        if idx is None:
            return []

        preds = []
        for pred_idx in self._graph.predecessor_indices(idx):
            edge_data: ToolEdgeData = self._graph.get_edge_data(pred_idx, idx)
            pred_name = self._idx_to_name.get(pred_idx, "?")
            preds.append((pred_name, edge_data.co_occurrence_count))

        preds.sort(key=lambda x: x[1], reverse=True)
        return preds[:top_n]

    def get_successors(self, tool_name: str, top_n: int = 3) -> List[Tuple[str, int]]:
        """Get top-N successors by co-occurrence count.

        Returns:
            List of (tool_name, count) tuples sorted by count descending.
        """
        idx = self._name_to_idx.get(tool_name)
        if idx is None:
            return []

        succs = []
        for succ_idx in self._graph.successor_indices(idx):
            edge_data: ToolEdgeData = self._graph.get_edge_data(idx, succ_idx)
            succ_name = self._idx_to_name.get(succ_idx, "?")
            succs.append((succ_name, edge_data.co_occurrence_count))

        succs.sort(key=lambda x: x[1], reverse=True)
        return succs[:top_n]

    def get_typical_chain(self, tool_name: str, length: int = 3) -> List[str]:
        """Get a typical workflow chain starting from (or including) tool_name.

        Follows the most frequent successor edges greedily.
        """
        chain = [tool_name]
        current = tool_name
        visited = {tool_name}

        for _ in range(length - 1):
            succs = self.get_successors(current, top_n=1)
            if not succs:
                break
            next_tool = succs[0][0]
            if next_tool in visited:
                break
            chain.append(next_tool)
            visited.add(next_tool)
            current = next_tool

        return chain

    # =========================================================================
    # RELATIONSHIP TEXT GENERATION
    # =========================================================================

    def get_relationship_text(
        self,
        tool_name: str,
        user_email: str = "",
        session_id: str = "",
    ) -> str:
        """
        Generate relationship text for embedding in the 'relationships' vector.

        Template:
            {tool_name} belongs to {service} service.
            Predecessors: {top_3}.
            Successors: {top_3}.
            User: {user_email}. Session: {session_id}.
            Core: {k_core}. Centrality: {betweenness:.2f}.
            Workflow: [{typical_chain}].
        """
        idx = self._name_to_idx.get(tool_name)
        if idx is None:
            return f"{tool_name} tool_response (no graph data)"

        node_data: ToolNodeData = self._graph.get_node_data(idx)

        # Ensure metrics are computed
        if self._metrics_stale:
            self.recompute_metrics()

        # Predecessors and successors
        preds = self.get_predecessors(tool_name)
        succs = self.get_successors(tool_name)
        pred_str = ", ".join(f"{n}({c})" for n, c in preds) if preds else "none"
        succ_str = ", ".join(f"{n}({c})" for n, c in succs) if succs else "none"

        # Metrics (rustworkx returns CentralityMapping / dict-like objects
        # that may not support .get(), so use index access with fallback)
        try:
            betweenness = self._betweenness[idx]
        except (KeyError, IndexError, TypeError):
            betweenness = 0.0
        try:
            k_core = self._k_cores[idx]
        except (KeyError, IndexError, TypeError):
            k_core = 0

        # Typical chain
        chain = self.get_typical_chain(tool_name)
        chain_str = " -> ".join(chain)

        parts = [
            f"{tool_name} belongs to {node_data.service} service.",
            f"Predecessors: {pred_str}.",
            f"Successors: {succ_str}.",
        ]
        if user_email:
            parts.append(f"User: {user_email}.")
        if session_id:
            parts.append(f"Session: {session_id}.")
        parts.append(f"Core: {k_core}. Centrality: {betweenness:.2f}.")
        parts.append(f"Workflow: [{chain_str}].")

        return " ".join(parts)

    # =========================================================================
    # SERIALIZATION
    # =========================================================================

    def get_stats(self) -> Dict[str, Any]:
        """Get graph statistics."""
        return {
            "nodes": self._graph.num_nodes(),
            "edges": self._graph.num_edges(),
            "sessions_tracked": len(self._session_history),
            "metrics_stale": self._metrics_stale,
        }

    def get_node_info(self, tool_name: str) -> Optional[Dict[str, Any]]:
        """Get info about a specific tool node."""
        idx = self._name_to_idx.get(tool_name)
        if idx is None:
            return None

        data: ToolNodeData = self._graph.get_node_data(idx)
        return {
            "tool_name": data.tool_name,
            "service": data.service,
            "call_count": data.call_count,
            "avg_execution_ms": data.avg_execution_ms,
            "in_degree": self._graph.in_degree(idx),
            "out_degree": self._graph.out_degree(idx),
            "betweenness": self._betweenness.get(idx, 0.0),
            "k_core": self._k_cores.get(idx, 0),
        }

    def clear_session(self, session_id: str) -> None:
        """Remove session history (called on session end)."""
        self._session_history.pop(session_id, None)

    @property
    def tool_names(self) -> Set[str]:
        """All tracked tool names."""
        return set(self._name_to_idx.keys())


__all__ = [
    "ToolRelationshipGraph",
    "ToolNodeData",
    "ToolEdgeData",
]
