# Familiar - Self-Hosted AI Agent Framework
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
ToolRouter — Guidance-Driven Tool Selection

Pre-filters the tool list before sending to the LLM, reducing token cost
by 60-80% and improving selection accuracy. Uses four scoring layers:

  1. Intent Match (0.4)     — TOOLS.yaml keyword patterns
  2. Chain Context (0.3)    — Boost chain targets of recent tool calls
  3. Skill Priority (0.2)   — TOOLS.yaml priority ranking
  4. Recency Bias (0.1)     — Recently-used tools get a small boost

The router never blocks valid tool calls. If confidence is low, it falls
back to sending more tools. If the LLM requests a tool not in the routed
set, the agent retries with the full set.

Usage:
    router = ToolRouter(guidance, tool_registry)
    result = router.route("check my email for urgent messages", context)
    # result.tools = ["check_inbox", "triage_inbox", "search_emails"]
    # result.schemas = [<3 schemas instead of 112>]
"""

from __future__ import annotations

import logging
import re
import threading
import time
from collections import defaultdict, deque
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set

if TYPE_CHECKING:
    from .guidance.loader import GuidanceLoader
    from .tools import ToolRegistry

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------


@dataclass
class RouterConfig:
    """Configuration for the ToolRouter."""

    # Maximum tools to send to LLM in the optimistic case
    max_tools: int = 8

    # Minimum composite score to include a tool
    confidence_threshold: float = 0.15

    # Number of tools to send when no confident matches (fallback)
    fallback_count: int = 15

    # Overall confidence below which we use fallback set
    min_routing_confidence: float = 0.20

    # Scoring weights (must sum to 1.0)
    weight_intent: float = 0.40
    weight_chain: float = 0.30
    weight_priority: float = 0.20
    weight_recency: float = 0.10

    # Number of recent turns to consider for chain/recency scoring
    recency_window: int = 5

    # Decay factor per turn for recency scoring (0.0 = no decay, 1.0 = instant)
    recency_decay: float = 0.3

    # Always include these tools regardless of routing
    always_include: List[str] = field(
        default_factory=lambda: [
            "accomplish",  # Planner entry point — catches everything
        ]
    )

    # Enable/disable the router (for A/B testing)
    enabled: bool = True


# ---------------------------------------------------------------
# Routing Result
# ---------------------------------------------------------------


@dataclass
class RoutingResult:
    """Result of tool routing for a user message."""

    # Selected tool names, in score order
    tools: List[str]

    # Tool schemas ready for LLM (same format as ToolRegistry.get_schemas())
    schemas: List[dict]

    # Tool name → composite score breakdown
    scores: Dict[str, ToolScore]

    # Overall routing confidence (0.0 - 1.0)
    confidence: float

    # True if we fell back to the broader set
    fallback: bool = False

    # Optional context hint to inject into the LLM system prompt
    context_hint: str = ""

    # Timing
    route_time_ms: float = 0.0

    @property
    def tool_count(self) -> int:
        return len(self.tools)

    @property
    def top_tool(self) -> Optional[str]:
        return self.tools[0] if self.tools else None


@dataclass
class ToolScore:
    """Score breakdown for a single tool."""

    tool_name: str
    intent_score: float = 0.0
    chain_score: float = 0.0
    priority_score: float = 0.0
    recency_score: float = 0.0
    composite: float = 0.0
    matched_keywords: List[str] = field(default_factory=list)
    chain_source: str = ""  # Which recent tool triggered the chain boost

    def compute_composite(self, config: RouterConfig) -> float:
        self.composite = (
            self.intent_score * config.weight_intent
            + self.chain_score * config.weight_chain
            + self.priority_score * config.weight_priority
            + self.recency_score * config.weight_recency
        )
        return self.composite


@dataclass
class RoutingOutcome:
    """Record of a single routing decision and its outcome."""

    routed_tools: list
    used_tools: list
    confidence: float
    expanded: bool
    timestamp: float = 0.0


# ---------------------------------------------------------------
# ToolRouter
# ---------------------------------------------------------------


class ToolRouter:
    """
    Guidance-driven tool routing.

    Scores all tools for a given user message and returns only the
    most relevant subset for the LLM to choose from.
    """

    def __init__(
        self,
        guidance: "GuidanceLoader",
        tool_registry: "ToolRegistry",
        config: Optional[RouterConfig] = None,
    ):
        self.guidance = guidance
        self.tool_registry = tool_registry
        self.config = config or RouterConfig()

        # Pre-compute data structures for fast scoring
        self._skill_priorities: Dict[str, int] = {}
        self._max_priority: int = 1
        self._tool_to_skill: Dict[str, str] = {}
        self._chain_targets: Dict[str, Set[str]] = defaultdict(set)
        self._all_tool_names: Set[str] = set()
        self._outcomes: deque = deque(maxlen=200)
        self._outcomes_lock = threading.Lock()

        self._build_index()

    def _build_index(self) -> None:
        """Build lookup tables from guidance data and schema registry."""
        # Always build schema lookup — needed for fallback intent scoring
        self._schema_lookup: Dict[str, dict] = {
            s["name"]: s for s in self.tool_registry.get_schemas()
        }

        if not self.guidance or not getattr(self.guidance, "skill_guidance", None):
            # No guidance — populate all_tool_names from registry only
            self._all_tool_names = set(self._schema_lookup.keys())
            logger.debug(f"ToolRouter index (schema-only): {len(self._all_tool_names)} tools")
            return

        for skill_name, sg in self.guidance.skill_guidance.items():
            self._skill_priorities[skill_name] = sg.priority
            if sg.priority > self._max_priority:
                self._max_priority = sg.priority

            for tool_name, tg in sg.tools.items():
                self._tool_to_skill[tool_name] = skill_name
                self._all_tool_names.add(tool_name)

                # Index chain targets: source_tool → {target_tools}
                for chain in tg.parsed_chains:
                    self._chain_targets[tool_name].add(chain.tool)

        # Also include tools from registry that may not be in TOOLS.yaml
        for schema in self.tool_registry.get_schemas():
            self._all_tool_names.add(schema["name"])

        logger.debug(
            f"ToolRouter index: {len(self._all_tool_names)} tools, "
            f"{sum(len(v) for v in self._chain_targets.values())} chain edges"
        )

    # ---------------------------------------------------------------
    # Main Routing Method
    # ---------------------------------------------------------------

    def route(
        self,
        message: str,
        context: Optional[Dict[str, Any]] = None,
        recent_tools: Optional[List[str]] = None,
        allowed_schemas: Optional[List[dict]] = None,
    ) -> RoutingResult:
        """
        Route a user message to the most relevant tools.

        Args:
            message: The user's message text
            context: Optional tool context dict
            recent_tools: Tool names called in recent turns (newest first)
            allowed_schemas: Pre-filtered schemas (from _get_allowed_tools).
                             If provided, routing only selects from this set.

        Returns:
            RoutingResult with selected tools and schemas.
        """
        start = time.monotonic()

        # Refresh _schema_lookup lazily if the registry has grown since _build_index
        # ran (which happens when skills load after ToolRouter construction).
        current_registry_size = len(self.tool_registry.get_schemas())
        if current_registry_size != len(self._schema_lookup):
            self._schema_lookup = {s["name"]: s for s in self.tool_registry.get_schemas()}
            self._all_tool_names = set(self._schema_lookup.keys())
            logger.debug(f"ToolRouter: refreshed schema index ({current_registry_size} tools)")

        if not self.config.enabled:
            # Router explicitly disabled — pass through all tools
            schemas = allowed_schemas or self.tool_registry.get_schemas()
            return RoutingResult(
                tools=[s["name"] for s in schemas],
                schemas=schemas,
                scores={},
                confidence=0.0,
                fallback=True,
                route_time_ms=(time.monotonic() - start) * 1000,
            )

        recent_tools = recent_tools or []

        # Build the allowed tool set from schemas
        if allowed_schemas:
            allowed_names = {s["name"] for s in allowed_schemas}
            schema_lookup = {s["name"]: s for s in allowed_schemas}
        else:
            all_schemas = self.tool_registry.get_schemas()
            allowed_names = {s["name"] for s in all_schemas}
            schema_lookup = {s["name"]: s for s in all_schemas}

        # Score every allowed tool
        scores: Dict[str, ToolScore] = {}
        for tool_name in allowed_names:
            ts = ToolScore(tool_name=tool_name)
            ts.intent_score = self._score_intent(tool_name, message)
            ts.chain_score = self._score_chain(tool_name, recent_tools)
            ts.priority_score = self._score_priority(tool_name)
            ts.recency_score = self._score_recency(tool_name, recent_tools)
            ts.matched_keywords = self._get_matched_keywords(tool_name, message)
            ts.compute_composite(self.config)
            scores[tool_name] = ts

        # Sort by composite score descending
        ranked = sorted(scores.values(), key=lambda s: s.composite, reverse=True)

        # Select tools above confidence threshold, up to max_tools
        selected: List[str] = []
        for ts in ranked:
            if len(selected) >= self.config.max_tools:
                break
            if ts.composite >= self.config.confidence_threshold:
                selected.append(ts.tool_name)

        # Always include pinned tools (if they're in allowed set)
        for pinned in self.config.always_include:
            if pinned in allowed_names and pinned not in selected:
                selected.append(pinned)

        # Compute overall confidence
        top_scores = [scores[t].composite for t in selected] if selected else [0.0]
        confidence = top_scores[0] if top_scores else 0.0

        # Fallback: if confidence too low or no tools selected, use broader set
        is_fallback = False
        if confidence < self.config.min_routing_confidence or len(selected) < 2:
            is_fallback = True
            selected = [ts.tool_name for ts in ranked[: self.config.fallback_count]]
            logger.debug(
                f"ToolRouter fallback: confidence={confidence:.3f}, "
                f"sending top {len(selected)} tools"
            )

        # Build context hint for LLM
        context_hint = self._build_context_hint(selected, scores, recent_tools)

        # Build schema list in selection order
        schemas = [schema_lookup[t] for t in selected if t in schema_lookup]

        elapsed = (time.monotonic() - start) * 1000

        result = RoutingResult(
            tools=selected,
            schemas=schemas,
            scores=scores,
            confidence=confidence,
            fallback=is_fallback,
            context_hint=context_hint,
            route_time_ms=elapsed,
        )

        logger.info(
            f"ToolRouter: {len(allowed_names)} tools → {len(selected)} "
            f"(confidence={confidence:.3f}, {'fallback' if is_fallback else 'routed'}, "
            f"{elapsed:.1f}ms)"
        )

        return result

    # ---------------------------------------------------------------
    # Retry: called when LLM requests a tool not in the routed set
    # ---------------------------------------------------------------

    def expand_for_retry(
        self,
        missing_tool: str,
        previous_result: RoutingResult,
        allowed_schemas: Optional[List[dict]] = None,
    ) -> RoutingResult:
        """
        Expand the tool set because the LLM requested a tool not in the
        routed set. Returns a new RoutingResult with the full allowed set.

        This guarantees correctness: routing never permanently blocks a tool.
        """
        schemas = allowed_schemas or self.tool_registry.get_schemas()
        logger.info(
            f"ToolRouter: expanding for '{missing_tool}' — sending all {len(schemas)} tools"
        )
        return RoutingResult(
            tools=[s["name"] for s in schemas],
            schemas=schemas,
            scores=previous_result.scores,
            confidence=0.0,
            fallback=True,
            context_hint="",
            route_time_ms=0.0,
        )

    # ---------------------------------------------------------------
    # Scoring Layers
    # ---------------------------------------------------------------

    # Common English stop words — present in nearly every message, no signal value
    _STOP_WORDS: frozenset = frozenset(
        {
            "the",
            "a",
            "an",
            "is",
            "to",
            "in",
            "for",
            "of",
            "and",
            "or",
            "my",
            "me",
            "i",
            "it",
            "its",
            "with",
            "at",
            "by",
            "from",
            "on",
            "be",
            "can",
            "will",
            "do",
            "what",
            "how",
            "when",
            "where",
            "which",
            "who",
            "please",
            "help",
            "get",
            "show",
            "make",
            "tell",
            "give",
            "list",
            "that",
            "this",
            "are",
            "was",
            "were",
            "has",
            "have",
            "had",
            "some",
            "all",
            "any",
            "just",
            "not",
            "but",
            "so",
            "if",
            "then",
            "than",
            "up",
            "out",
            "about",
            "into",
            "after",
            "before",
            "over",
        }
    )

    def _score_intent_schema(self, tool_name: str, message: str) -> float:
        """
        Schema-based fallback intent scorer.

        When TOOLS.yaml guidance is unavailable, score each tool by matching
        the user's message words against the tool's name and description tokens.
        Tool-name matches are weighted 3× over description matches to reward
        precise lexical alignment (e.g. 'add task' → 'add_task').

        Returns a score in [0.0, 1.0].
        """
        schema = self._schema_lookup.get(tool_name)
        if not schema:
            return 0.0

        msg_words = set(re.findall(r"\w+", message.lower())) - self._STOP_WORDS
        if not msg_words:
            return 0.0

        name_tokens = (
            set(re.findall(r"\w+", schema.get("name", "").replace("_", " ").lower()))
            - self._STOP_WORDS
        )
        desc_tokens = (
            set(re.findall(r"\w+", schema.get("description", "").lower())) - self._STOP_WORDS
        )

        name_hits = len(msg_words & name_tokens)
        desc_hits = len(msg_words & desc_tokens)

        raw = name_hits * 3.0 + desc_hits * 1.0
        # Normalise: full name match of a k-word query scores 1.0
        norm = len(msg_words) * 3.0 * 0.4
        return min(1.0, raw / max(1.0, norm))

    def _score_intent(self, tool_name: str, message: str) -> float:
        """
        Layer 1: Intent matching.

        Primary path: TOOLS.yaml keyword patterns (when guidance loaded).
        Fallback path: schema-based token overlap (always available).
        """
        if self.guidance and getattr(self.guidance, "skill_guidance", None):
            # ── Primary: guidance-driven scoring ────────────────────────
            message_lower = message.lower()
            words = set(re.findall(r"\w+", message_lower))

            skill_name = self._tool_to_skill.get(tool_name)
            if not skill_name:
                return self._score_intent_schema(tool_name, message)

            sg = self.guidance.skill_guidance.get(skill_name)
            if not sg:
                return self._score_intent_schema(tool_name, message)

            max_score = 0.0
            for ip in sg.intent_patterns:
                if ip.suggest != tool_name:
                    continue
                score = 0
                for kw in ip.keywords:
                    kw_lower = kw.lower()
                    if kw_lower in words:
                        score += 1
                    elif len(kw_lower) > 3 and kw_lower in message_lower:
                        score += 0.7
                if score > 0:
                    normalized = min(1.0, score / max(1, len(ip.keywords) * 0.5))
                    max_score = max(max_score, normalized)

            tg = sg.tools.get(tool_name)
            if tg and tg.example_prompts:
                for prompt in tg.example_prompts:
                    prompt_words = set(re.findall(r"\w+", prompt.lower()))
                    overlap = len(words & prompt_words)
                    if overlap >= 2:
                        prompt_score = min(1.0, overlap / max(1, len(prompt_words) * 0.4))
                        max_score = max(max_score, prompt_score * 0.8)

            # Blend with schema score so partial guidance still benefits
            schema_score = self._score_intent_schema(tool_name, message)
            return max(max_score, schema_score * 0.6)

        # ── Fallback: schema-only scoring ───────────────────────────────
        return self._score_intent_schema(tool_name, message)

    def _score_chain(self, tool_name: str, recent_tools: List[str]) -> float:
        """
        Layer 2: Chain context — boost tools that are chain targets
        of recently-called tools.

        If check_inbox was called last turn, and check_inbox chains_with
        triage_inbox, then triage_inbox gets a high chain score.
        """
        if not recent_tools:
            return 0.0

        max_score = 0.0
        window = recent_tools[: self.config.recency_window]

        for i, recent_tool in enumerate(window):
            targets = self._chain_targets.get(recent_tool, set())
            if tool_name in targets:
                # Decay by position: most recent = 1.0, older = less
                decay = (1.0 - self.config.recency_decay) ** i
                score = decay
                max_score = max(max_score, score)

        return max_score

    def _score_priority(self, tool_name: str) -> float:
        """
        Layer 3: Skill priority from TOOLS.yaml.

        Normalized to 0.0 - 1.0 by dividing by max priority.
        """
        skill_name = self._tool_to_skill.get(tool_name)
        if not skill_name:
            return 0.0

        priority = self._skill_priorities.get(skill_name, 0)
        return priority / max(1, self._max_priority)

    def _score_recency(self, tool_name: str, recent_tools: List[str]) -> float:
        """
        Layer 4: Recency bias — tools used recently get a small boost
        for conversation continuity.
        """
        if not recent_tools:
            return 0.0

        window = recent_tools[: self.config.recency_window]
        for i, recent in enumerate(window):
            if recent == tool_name:
                return (1.0 - self.config.recency_decay) ** i
        return 0.0

    # ---------------------------------------------------------------
    # Helpers
    # ---------------------------------------------------------------

    def _get_matched_keywords(self, tool_name: str, message: str) -> List[str]:
        """Get which keywords matched for this tool (for diagnostics)."""
        skill_name = self._tool_to_skill.get(tool_name)
        if not skill_name:
            return []

        sg = self.guidance.skill_guidance.get(skill_name)
        if not sg:
            return []

        message_lower = message.lower()
        matched = []
        for ip in sg.intent_patterns:
            if ip.suggest == tool_name:
                for kw in ip.keywords:
                    if kw.lower() in message_lower:
                        matched.append(kw)
        return matched

    def _build_context_hint(
        self,
        selected: List[str],
        scores: Dict[str, ToolScore],
        recent_tools: List[str],
    ) -> str:
        """
        Build an optional system-level hint for the LLM explaining
        why these tools were pre-selected.

        Only generated when there's a clear chain context to convey.
        """
        # Only hint if there's chain context worth mentioning
        chain_hints = []
        for tool_name in selected:
            ts = scores.get(tool_name)
            if ts and ts.chain_score > 0.3 and recent_tools:
                # Find which recent tool triggered the chain
                for recent in recent_tools[:3]:
                    if tool_name in self._chain_targets.get(recent, set()):
                        chain_hints.append(f"{tool_name} (follows from {recent})")
                        break

        if chain_hints:
            return f"Based on conversation flow, relevant follow-up tools: {', '.join(chain_hints)}"
        return ""

    # ---------------------------------------------------------------
    # Diagnostics
    # ---------------------------------------------------------------

    def explain(self, message: str, recent_tools: Optional[List[str]] = None) -> str:
        """
        Human-readable explanation of how routing would work for a message.
        Useful for debugging and tuning.
        """
        result = self.route(message, recent_tools=recent_tools)

        lines = [
            f'Message: "{message}"',
            f"Routed: {result.tool_count} tools "
            f"(confidence={result.confidence:.3f}, "
            f"{'fallback' if result.fallback else 'routed'}, "
            f"{result.route_time_ms:.1f}ms)",
            "",
        ]

        for tool_name in result.tools[:10]:
            ts = result.scores.get(tool_name)
            if ts:
                parts = []
                if ts.intent_score > 0:
                    kws = ", ".join(ts.matched_keywords) if ts.matched_keywords else "example"
                    parts.append(f"intent={ts.intent_score:.2f} [{kws}]")
                if ts.chain_score > 0:
                    parts.append(f"chain={ts.chain_score:.2f}")
                if ts.priority_score > 0:
                    parts.append(f"priority={ts.priority_score:.2f}")
                if ts.recency_score > 0:
                    parts.append(f"recency={ts.recency_score:.2f}")
                breakdown = " + ".join(parts) if parts else "no signals"
                lines.append(f"  {ts.composite:.3f}  {tool_name}  ({breakdown})")

        return "\n".join(lines)

    def record_routing_outcome(
        self,
        routed_tools: List[str],
        used_tools: List[str],
        confidence: float,
        expanded: bool = False,
    ) -> None:
        """Record the outcome of a routing decision for feedback analysis."""
        with self._outcomes_lock:
            self._outcomes.append(
                RoutingOutcome(
                    routed_tools=routed_tools,
                    used_tools=used_tools,
                    confidence=confidence,
                    expanded=expanded,
                    timestamp=time.monotonic(),
                )
            )

    def get_routing_stats(self) -> dict:
        """Return aggregate stats over recorded routing outcomes."""
        with self._outcomes_lock:
            snapshot = list(self._outcomes)
        total = len(snapshot)
        if total == 0:
            return {
                "total": 0,
                "hit_rate": 0.0,
                "expansion_rate": 0.0,
                "avg_confidence": 0.0,
            }
        hits = sum(1 for o in snapshot if set(o.used_tools).issubset(set(o.routed_tools)))
        expansions = sum(1 for o in snapshot if o.expanded)
        avg_conf = sum(o.confidence for o in snapshot) / total
        return {
            "total": total,
            "hit_rate": hits / total,
            "expansion_rate": expansions / total,
            "avg_confidence": avg_conf,
        }
