"""
Semantic router for Agent-Gantry.

Intelligent tool selection using semantic search, intent classification, and context.
"""

from __future__ import annotations

import math
import re
from dataclasses import dataclass
from enum import Enum
from functools import lru_cache
from time import perf_counter
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from agent_gantry.adapters.embedders.base import EmbeddingAdapter
    from agent_gantry.adapters.llm_client import LLMClient
    from agent_gantry.adapters.rerankers.base import RerankerAdapter
    from agent_gantry.adapters.vector_stores.base import VectorStoreAdapter
    from agent_gantry.schema.query import ToolQuery
    from agent_gantry.schema.tool import ToolDefinition


@lru_cache(maxsize=256)
def _get_token_pattern(token: str) -> re.Pattern[str]:
    """Cache compiled regex patterns for token matching."""
    return re.compile(rf"(?<!\w){re.escape(token)}(?!\w)")


class TaskIntent(str, Enum):
    """High-level task intent classification."""

    DATA_QUERY = "data_query"
    DATA_MUTATION = "data_mutation"
    ANALYSIS = "analysis"
    COMMUNICATION = "communication"
    FILE_OPERATIONS = "file_operations"
    CUSTOMER_SUPPORT = "customer_support"
    ADMIN = "admin"
    UNKNOWN = "unknown"


INTENT_TAG_MAPPING: dict[TaskIntent, list[str]] = {
    TaskIntent.DATA_QUERY: ["query", "search", "get", "list", "fetch", "read"],
    TaskIntent.DATA_MUTATION: ["create", "update", "delete", "write", "modify"],
    TaskIntent.ANALYSIS: ["analyze", "compute", "aggregate", "calculate", "report"],
    TaskIntent.COMMUNICATION: ["email", "message", "notify", "send", "chat", "dm", "teams", "discord", "communicate"],
    TaskIntent.FILE_OPERATIONS: ["file", "upload", "download", "convert", "export"],
    TaskIntent.CUSTOMER_SUPPORT: ["ticket", "refund", "support", "customer"],
    TaskIntent.ADMIN: ["user", "permission", "setting", "config", "admin"],
}


@dataclass
class RoutingSignals:
    """Signals used for tool scoring."""

    semantic_similarity: float
    intent_match: float
    conversation_relevance: float
    health_score: float
    cost_score: float
    already_used_penalty: float
    already_failed_penalty: float
    deprecated_penalty: float


@dataclass
class RoutingWeights:
    """Weights for combining routing signals."""

    semantic: float = 0.6
    intent: float = 0.15
    conversation: float = 0.1
    health: float = 0.1
    cost: float = 0.05


@dataclass
class RoutingResult:
    """Routing outcome with timing metadata."""

    tools: list[tuple[ToolDefinition, float]]
    query_embedding_time_ms: float
    vector_search_time_ms: float
    rerank_time_ms: float | None
    candidate_count: int
    filtered_count: int


def compute_final_score(signals: RoutingSignals, weights: RoutingWeights) -> float:
    """
    Compute the final score for a tool.

    Args:
        signals: The routing signals for the tool
        weights: The weights for each signal

    Returns:
        The final composite score
    """
    base_score = (
        signals.semantic_similarity * weights.semantic
        + signals.intent_match * weights.intent
        + signals.conversation_relevance * weights.conversation
        + signals.health_score * weights.health
        + signals.cost_score * weights.cost
    )
    penalties = (
        signals.already_used_penalty
        + signals.already_failed_penalty
        + signals.deprecated_penalty
    )
    return max(0.0, base_score - penalties)


async def classify_intent(
    query: str,
    conversation_summary: str | None = None,
    use_llm: bool = False,
    llm_client: LLMClient | None = None,
) -> TaskIntent:
    """
    Classify the intent of a query.

    Args:
        query: The user's query
        conversation_summary: Optional conversation context
        use_llm: Whether to use LLM for classification
        llm_client: Optional LLM client for classification

    Returns:
        The classified intent
    """
    query_lower = query.lower()
    scores: dict[TaskIntent, int] = {}

    enriched_query = query_lower
    if conversation_summary:
        enriched_query = f"{enriched_query} {conversation_summary.lower()}"

    # First try keyword-based classification
    for intent, keywords in INTENT_TAG_MAPPING.items():
        scores[intent] = sum(1 for kw in keywords if kw in enriched_query)

    if max(scores.values()) > 0:
        return max(scores, key=lambda k: scores[k])

    # Fall back to LLM-based classification if enabled and available
    if use_llm and llm_client:
        available_intents = [intent.value for intent in TaskIntent]
        try:
            result = await llm_client.classify_intent(
                query=query,
                conversation_summary=conversation_summary,
                available_intents=available_intents,
            )
            # Convert string result to TaskIntent enum
            for intent in TaskIntent:
                if intent.value == result:
                    return intent
        except Exception:
            # Fall back to UNKNOWN if LLM classification fails
            pass

    return TaskIntent.UNKNOWN


class SemanticRouter:
    """
    Semantic router for intelligent tool selection.

    Uses:
    - Vector similarity search
    - Intent classification
    - Conversation context
    - Tool health metrics
    - MMR diversity
    """

    def __init__(
        self,
        vector_store: VectorStoreAdapter,
        embedder: EmbeddingAdapter,
        reranker: RerankerAdapter | None = None,
        weights: RoutingWeights | None = None,
        llm_client: LLMClient | None = None,
        use_llm_for_intent: bool = False,
    ) -> None:
        """
        Initialize the semantic router.

        Args:
            vector_store: Vector store for tool embeddings
            embedder: Embedding model for queries
            reranker: Optional reranker for precision
            weights: Routing signal weights
            llm_client: Optional LLM client for intent classification
            use_llm_for_intent: Whether to use LLM for intent classification
        """
        self._vector_store = vector_store
        self._embedder = embedder
        self._reranker = reranker
        self._weights = weights or RoutingWeights()
        self._llm_client = llm_client
        self._use_llm_for_intent = use_llm_for_intent

    async def route(
        self,
        query: ToolQuery,
    ) -> RoutingResult:
        """
        Route a query to the most relevant tools.

        Args:
            query: The tool query with context

        Returns:
            Routing result with scores and timings
        """
        embed_start = perf_counter()
        query_embedding = await self._embedder.embed_text(query.context.query)
        query_embedding_time_ms = (perf_counter() - embed_start) * 1000

        filters: dict[str, list[str]] | None = None
        if query.namespaces:
            filters = {"namespace": query.namespaces}

        # Request embeddings if MMR will be used (optimization to avoid re-embedding)
        include_embeddings = query.diversity_factor > 0

        search_start = perf_counter()
        candidates = await self._vector_store.search(
            query_vector=query_embedding,
            limit=query.limit * 4,
            filters=filters,
            score_threshold=query.score_threshold,
            include_embeddings=include_embeddings,
        )
        vector_search_time_ms = (perf_counter() - search_start) * 1000

        intent = await self._resolve_intent(query)

        # Store embeddings separately for MMR if they were included
        tool_embeddings: dict[str, list[float]] = {}
        scored_tools: list[tuple[ToolDefinition, float]] = []

        for candidate in candidates:
            if include_embeddings:
                tool, semantic_score, embedding = candidate  # type: ignore
                tool_key = f"{tool.namespace}.{tool.name}"
                tool_embeddings[tool_key] = embedding
            else:
                tool, semantic_score = candidate  # type: ignore
            if query.exclude_deprecated and tool.deprecated:
                continue
            if query.namespaces and tool.namespace not in query.namespaces:
                continue
            if query.required_capabilities and not all(
                cap in tool.capabilities for cap in query.required_capabilities
            ):
                continue
            if query.excluded_capabilities and any(
                cap in tool.capabilities for cap in query.excluded_capabilities
            ):
                continue
            if query.sources and tool.source not in query.sources:
                continue
            if query.exclude_unhealthy and tool.health.circuit_breaker_open:
                continue

            signals = self._compute_signals(
                tool=tool,
                semantic_score=semantic_score,
                intent=intent,
                query=query,
            )
            final_score = compute_final_score(signals, self._weights)
            scored_tools.append((tool, final_score))

        scored_tools.sort(key=lambda x: x[1], reverse=True)

        rerank_time_ms: float | None = None
        if self._reranker and query.enable_reranking:
            rerank_start = perf_counter()
            scored_tools = await self._reranker.rerank(
                query.context.query,
                scored_tools,
                query.limit,
            )
            rerank_time_ms = (perf_counter() - rerank_start) * 1000

        if query.diversity_factor > 0 and len(scored_tools) > 1:
            scored_tools = await self._apply_mmr(
                scored_tools,
                query_embedding,
                query.diversity_factor,
                query.limit,
                tool_embeddings,
            )

        final_tools = scored_tools[: query.limit]
        return RoutingResult(
            tools=final_tools,
            query_embedding_time_ms=query_embedding_time_ms,
            vector_search_time_ms=vector_search_time_ms,
            rerank_time_ms=rerank_time_ms,
            candidate_count=len(candidates),
            filtered_count=len(scored_tools),
        )

    async def _resolve_intent(self, query: ToolQuery) -> TaskIntent:
        """Resolve intent using context override or classification."""
        if query.context.inferred_intent:
            try:
                return TaskIntent(query.context.inferred_intent)
            except ValueError:
                pass
        return await classify_intent(
            query.context.query,
            query.context.conversation_summary,
            use_llm=self._use_llm_for_intent,
            llm_client=self._llm_client,
        )

    def _compute_signals(
        self,
        tool: ToolDefinition,
        semantic_score: float,
        intent: TaskIntent,
        query: ToolQuery,
    ) -> RoutingSignals:
        """Compute routing signals for a tool."""
        # Intent match
        intent_match = 0.0
        if intent != TaskIntent.UNKNOWN:
            intent_keywords = INTENT_TAG_MAPPING.get(intent, [])
            tool_text = f"{tool.name.lower()} {tool.description.lower()} {' '.join(tool.tags).lower()}"
            if any(kw in tool_text for kw in intent_keywords):
                intent_match = 1.0

        # Conversation relevance
        conversation_relevance = 0.0
        if tool.name in query.context.tools_already_used:
            conversation_relevance = 0.5
        summary = (query.context.conversation_summary or "").lower()
        if summary and self._contains_token(summary, tool.name.lower()):
            conversation_relevance = min(1.0, conversation_relevance + 0.2)
        for message in query.context.recent_messages:
            content = message.get("content", "").lower()
            if content and self._contains_token(content, tool.name.lower()):
                conversation_relevance = min(1.0, conversation_relevance + 0.2)

        # Health score
        health_score = 0.0 if tool.health.circuit_breaker_open else tool.health.success_rate

        # Cost score (inverse - lower cost is better)
        cost_score = 1.0 - min(tool.cost.estimated_latency_ms / 10000, 1.0)

        # Penalties
        already_used_penalty = 0.1 if tool.name in query.context.tools_already_used else 0.0
        already_failed_penalty = 0.3 if tool.name in query.context.tools_failed else 0.0
        deprecated_penalty = 0.5 if tool.deprecated else 0.0

        return RoutingSignals(
            semantic_similarity=semantic_score,
            intent_match=intent_match,
            conversation_relevance=conversation_relevance,
            health_score=health_score,
            cost_score=cost_score,
            already_used_penalty=already_used_penalty,
            already_failed_penalty=already_failed_penalty,
            deprecated_penalty=deprecated_penalty,
        )

    async def _apply_mmr(
        self,
        scored_tools: list[tuple[ToolDefinition, float]],
        query_embedding: list[float],
        diversity_factor: float,
        limit: int,
        cached_embeddings: dict[str, list[float]] | None = None,
    ) -> list[tuple[ToolDefinition, float]]:
        """Apply MMR to encourage diversity in selected tools."""
        if not scored_tools:
            return []

        _ = query_embedding  # kept for signature consistency; not used in relevance scoring
        lambda_param = 1.0 - diversity_factor
        relevance_scores = [score for _, score in scored_tools]

        # Use cached embeddings if available, otherwise re-embed (fallback)
        if cached_embeddings:
            embeddings = []
            for tool, _ in scored_tools:
                tool_key = f"{tool.namespace}.{tool.name}"
                embedding = cached_embeddings.get(tool_key)
                if embedding is None:
                    # Fallback: re-embed this specific tool if not in cache
                    embedding = await self._embedder.embed_text(tool.to_searchable_text())
                embeddings.append(embedding)
        else:
            # Fallback: re-embed all tools (old behavior for backward compatibility)
            tool_texts = [tool.to_searchable_text() for tool, _ in scored_tools]
            embeddings = await self._embedder.embed_batch(tool_texts)

        selected: list[int] = []
        candidates = list(range(len(scored_tools)))

        first_idx = max(candidates, key=lambda i: relevance_scores[i])
        selected.append(first_idx)
        candidates.remove(first_idx)

        while candidates and len(selected) < limit:
            mmr_scores: dict[int, float] = {}
            for idx in candidates:
                similarity_to_selected = max(
                    (self._cosine_similarity(embeddings[idx], embeddings[sel]) for sel in selected),
                    default=0.0,
                )
                mmr_scores[idx] = lambda_param * relevance_scores[idx] - (
                    1.0 - lambda_param
                ) * similarity_to_selected

            next_idx = max(mmr_scores, key=lambda k: mmr_scores[k])
            selected.append(next_idx)
            candidates.remove(next_idx)

        return [scored_tools[i] for i in selected]

    def _cosine_similarity(self, a: list[float], b: list[float]) -> float:
        """Calculate cosine similarity between two vectors."""
        if not a or not b or len(a) != len(b):
            return 0.0
        dot_product = sum(x * y for x, y in zip(a, b))
        norm_a = math.sqrt(sum(x * x for x in a))
        norm_b = math.sqrt(sum(y * y for y in b))
        if norm_a == 0.0 or norm_b == 0.0:
            return 0.0
        return dot_product / (norm_a * norm_b)

    def _contains_token(self, text: str, token: str) -> bool:
        """Return True if token appears as a standalone word in text."""
        if not token:
            return False
        return _get_token_pattern(token).search(text) is not None
