"""Query Router for SAGE Engine.

Classifies incoming queries and determines the appropriate processing path.
Routes queries to knowledge (RAG), calculation (tools), document generation,
analysis, or training handlers.
"""

import re
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

from sage.engine.config import RouterConfig


class QueryType(Enum):
    """Types of queries that SAGE can handle."""

    KNOWLEDGE = "knowledge"  # Factual Q&A, definitions, explanations
    CALCULATION = "calculation"  # Needs tool support (k-eff, unit conversion)
    DOCUMENT = "document"  # CSE drafting, technical writing
    TRAINING = "training"  # Educational, tutorial content
    ANALYSIS = "analysis"  # Complex multi-step reasoning


class ComplexityLevel(Enum):
    """Estimated complexity of a query."""

    SIMPLE = "simple"  # Single-step, direct answer
    MODERATE = "moderate"  # Multi-source, some reasoning
    COMPLEX = "complex"  # Multi-step, tool integration, verification


@dataclass
class RoutingDecision:
    """Result of query routing analysis."""

    query_type: QueryType
    complexity: ComplexityLevel
    confidence: float  # Confidence in the routing decision
    required_tools: list[str] = field(default_factory=list)
    tool_params: dict[str, Any] = field(default_factory=dict)
    reasoning_steps: int = 1
    requires_rag: bool = True
    requires_verification: bool = False
    routing_reasons: list[str] = field(default_factory=list)


class QueryRouter:
    """Routes queries to appropriate handlers based on classification.

    The router analyzes incoming queries to determine:
    1. Query type (knowledge, calculation, document, training, analysis)
    2. Required tools for the query
    3. Complexity level and reasoning steps needed
    4. Whether RAG context is needed
    """

    def __init__(self, config: RouterConfig | None = None):
        """Initialize the query router.

        Args:
            config: Router configuration. Uses defaults if not provided.
        """
        self.config = config or RouterConfig()

        # Pre-compile keyword patterns for efficiency
        self._calculation_pattern = self._compile_keywords(self.config.calculation_keywords)
        self._knowledge_pattern = self._compile_keywords(self.config.knowledge_keywords)
        self._document_pattern = self._compile_keywords(self.config.document_keywords)
        self._analysis_pattern = self._compile_keywords(self.config.analysis_keywords)
        self._training_pattern = self._compile_keywords(self.config.training_keywords)

        # Tool detection patterns
        self._tool_patterns = {
            "keff_estimator": re.compile(
                r"\b(k-?eff|criticality|critical\s*mass|subcritical|supercritical)\b",
                re.IGNORECASE,
            ),
            "unit_converter": re.compile(
                r"\b(convert|conversion|grams?\s+to|kg\s+to|liters?\s+to|"
                r"enrichment|wt\s*%|atom\s*%)\b",
                re.IGNORECASE,
            ),
            "scale_interface": re.compile(
                r"\b(scale|keno|csas|origen|triton|input\s+file|run\s+calculation)\b",
                re.IGNORECASE,
            ),
            "standards_lookup": re.compile(
                r"\b(ansi|ans-?\d|10\s*cfr|nureg|regulatory\s+guide|standard)\b",
                re.IGNORECASE,
            ),
            "geometry_visualizer": re.compile(
                r"\b(visualize|visualization|diagram|draw|show|display|svg|"
                r"cross[\s-]?section|geometry\s+(image|picture|diagram))\b",
                re.IGNORECASE,
            ),
        }

    def _compile_keywords(self, keywords: list[str]) -> re.Pattern:
        """Compile keywords into a regex pattern.

        Args:
            keywords: List of keyword strings.

        Returns:
            Compiled regex pattern.
        """
        escaped = [re.escape(kw) for kw in keywords]
        pattern = r"\b(" + "|".join(escaped) + r")\b"
        return re.compile(pattern, re.IGNORECASE)

    def route(self, query: str, context: dict[str, Any] | None = None) -> RoutingDecision:
        """Route a query to the appropriate handler.

        Args:
            query: The user's query string.
            context: Optional context dictionary with additional info.

        Returns:
            RoutingDecision with query type, tools, and processing params.
        """
        context = context or {}

        # Classify query type
        query_type = self.classify_query_type(query)

        # Identify required tools
        required_tools = self.identify_required_tools(query)

        # Estimate complexity
        complexity = self.estimate_complexity(query, query_type, required_tools)

        # Build routing decision
        decision = RoutingDecision(
            query_type=query_type,
            complexity=complexity,
            confidence=self._calculate_routing_confidence(query, query_type),
            required_tools=required_tools,
            tool_params=self._extract_tool_params(query, required_tools),
            reasoning_steps=self._estimate_reasoning_steps(complexity, query_type),
            requires_rag=self._needs_rag(query_type, required_tools),
            requires_verification=self._needs_verification(query_type, complexity),
            routing_reasons=self._get_routing_reasons(query, query_type, required_tools),
        )

        # Override based on context flags
        if context.get("force_knowledge_only"):
            decision.query_type = QueryType.KNOWLEDGE
            decision.required_tools = []

        if context.get("allow_tools") is False:
            decision.required_tools = []

        return decision

    def classify_query_type(self, query: str) -> QueryType:
        """Classify the type of query.

        Args:
            query: The user's query string.

        Returns:
            The classified QueryType.
        """
        query_lower = query.lower()

        # Score each type based on keyword matches
        scores: dict[QueryType, int] = {
            QueryType.KNOWLEDGE: 0,
            QueryType.CALCULATION: 0,
            QueryType.DOCUMENT: 0,
            QueryType.TRAINING: 0,
            QueryType.ANALYSIS: 0,
        }

        # Check for strong intent indicators at the START of query
        # These take priority over domain keywords in the middle
        knowledge_start_patterns = [
            r"^what\s+is\b",
            r"^define\b",
            r"^explain\b",
            r"^describe\b",
            r"^who\b",
            r"^when\b",
            r"^where\b",
            r"^why\b",
        ]
        training_start_patterns = [
            r"^teach\b",
            r"^walk\s+me\s+through\b",
            r"^how\s+do\s+i\b",
            r"^show\s+me\s+how\b",
            r"^learn\b",
        ]
        calculation_start_patterns = [
            r"^calculate\b",
            r"^compute\b",
            r"^estimate\b",
            r"^convert\b",
            r"^check\s+if\b",
        ]

        # Strong start pattern bonus
        has_knowledge_start = any(
            re.search(p, query_lower) for p in knowledge_start_patterns
        )
        has_training_start = any(
            re.search(p, query_lower) for p in training_start_patterns
        )
        has_calculation_start = any(
            re.search(p, query_lower) for p in calculation_start_patterns
        )

        if has_knowledge_start:
            scores[QueryType.KNOWLEDGE] += 4
        if has_training_start:
            scores[QueryType.TRAINING] += 4
        if has_calculation_start:
            scores[QueryType.CALCULATION] += 4

        # Check patterns and accumulate scores
        if self._calculation_pattern.search(query_lower):
            scores[QueryType.CALCULATION] += 3

        if self._knowledge_pattern.search(query_lower):
            scores[QueryType.KNOWLEDGE] += 2

        if self._document_pattern.search(query_lower):
            scores[QueryType.DOCUMENT] += 3

        if self._analysis_pattern.search(query_lower):
            scores[QueryType.ANALYSIS] += 2

        if self._training_pattern.search(query_lower):
            scores[QueryType.TRAINING] += 2

        # Tool keywords suggest calculation, but only if no strong
        # knowledge/training start indicators are present
        if not has_knowledge_start and not has_training_start:
            for tool, pattern in self._tool_patterns.items():
                if pattern.search(query_lower):
                    if tool in ("keff_estimator", "unit_converter", "scale_interface"):
                        scores[QueryType.CALCULATION] += 2

        # Numeric values suggest calculation
        if re.search(r"\d+\.?\d*\s*(kg|g|mol|L|cm|m|%)", query):
            scores[QueryType.CALCULATION] += 1

        # Mathematical expressions and equations suggest calculation
        # Patterns: "x = ?", "solve for", mathematical operators
        if re.search(r"[=+\-*/^]\s*\?", query):  # Equation with unknown
            scores[QueryType.CALCULATION] += 2
        if re.search(r"\b(x|y|k)\s*=", query_lower):  # Variable assignment
            scores[QueryType.CALCULATION] += 1
        if re.search(r"[+\-*/]\s*\d", query):  # Arithmetic operations
            scores[QueryType.CALCULATION] += 1

        # Questions with "?" default toward knowledge
        if query.strip().endswith("?"):
            scores[QueryType.KNOWLEDGE] += 1

        # Return highest scoring type, defaulting to KNOWLEDGE
        max_score = max(scores.values())
        if max_score == 0:
            return QueryType.KNOWLEDGE

        for query_type, score in scores.items():
            if score == max_score:
                return query_type

        return QueryType.KNOWLEDGE

    def identify_required_tools(self, query: str) -> list[str]:
        """Identify which tools are needed for the query.

        Args:
            query: The user's query string.

        Returns:
            List of tool names required.
        """
        required = []
        query_lower = query.lower()

        # Detect if this is a conceptual/educational query where tools aren't needed
        # even if domain keywords are present
        conceptual_patterns = [
            r"^what\s+is\b(?!.*\bin\b)",  # "what is X" but NOT "what is X in Y" (conversion)
            r"^define\b",
            r"^explain\b",
            r"^describe\b",
            r"^teach\b",
            r"^walk\s+me\s+through\b",
            r"^how\s+do\s+i\b",
            r"^learn\b",
            r"\bbasics\b",
            r"\bintroduction\b",
            r"\btutorial\b",
        ]
        is_conceptual = any(re.search(p, query_lower) for p in conceptual_patterns)

        # "What is X in Y" pattern is a conversion request, not conceptual
        is_conversion_query = re.search(r"^what\s+is\b.*\bin\b", query_lower) is not None

        # For conceptual queries, only match tools if there's an explicit action
        # word suggesting actual calculation/lookup is needed
        action_words = ["calculate", "compute", "estimate", "convert", "check if", "run"]
        has_action = any(word in query_lower for word in action_words)

        for tool_name, pattern in self._tool_patterns.items():
            if pattern.search(query_lower):
                # Skip tool identification for conceptual queries without action words
                # But allow conversion queries ("what is X in Y")
                if is_conceptual and not has_action and not is_conversion_query:
                    continue
                required.append(tool_name)

        return required

    def estimate_complexity(
        self,
        query: str,
        query_type: QueryType,
        required_tools: list[str],
    ) -> ComplexityLevel:
        """Estimate the complexity of processing a query.

        Args:
            query: The user's query string.
            query_type: Classified query type.
            required_tools: List of required tools.

        Returns:
            Estimated complexity level.
        """
        complexity_score = 0

        # Tool usage increases complexity
        complexity_score += len(required_tools) * 2

        # Query type affects complexity
        if query_type == QueryType.ANALYSIS:
            complexity_score += 3
        elif query_type == QueryType.DOCUMENT:
            complexity_score += 2
        elif query_type == QueryType.CALCULATION:
            complexity_score += 1

        # Multiple questions increase complexity
        question_count = query.count("?")
        if question_count > 1:
            complexity_score += question_count

        # Conditional language suggests analysis
        if re.search(r"\b(if|when|whether|assuming|given that)\b", query, re.IGNORECASE):
            complexity_score += 1

        # Comparative language suggests analysis
        if re.search(r"\b(compare|versus|vs\.?|better|worse|difference)\b", query, re.IGNORECASE):
            complexity_score += 1

        # Map score to complexity level
        if complexity_score <= 1:
            return ComplexityLevel.SIMPLE
        elif complexity_score <= 4:
            return ComplexityLevel.MODERATE
        else:
            return ComplexityLevel.COMPLEX

    def _calculate_routing_confidence(self, query: str, query_type: QueryType) -> float:
        """Calculate confidence in the routing decision.

        Args:
            query: The user's query string.
            query_type: Classified query type.

        Returns:
            Confidence score between 0 and 1.
        """
        # Start with base confidence
        confidence = 0.7

        # Strong keyword matches increase confidence
        query_lower = query.lower()

        pattern_map = {
            QueryType.CALCULATION: self._calculation_pattern,
            QueryType.KNOWLEDGE: self._knowledge_pattern,
            QueryType.DOCUMENT: self._document_pattern,
            QueryType.ANALYSIS: self._analysis_pattern,
            QueryType.TRAINING: self._training_pattern,
        }

        pattern = pattern_map.get(query_type)
        if pattern:
            matches = pattern.findall(query_lower)
            confidence += min(len(matches) * 0.1, 0.2)

        # Tool patterns for calculation queries
        if query_type == QueryType.CALCULATION:
            tool_matches = sum(
                1 for p in self._tool_patterns.values() if p.search(query_lower)
            )
            confidence += min(tool_matches * 0.05, 0.1)

        return min(confidence, 1.0)

    def _extract_tool_params(
        self, query: str, required_tools: list[str]
    ) -> dict[str, Any]:
        """Extract parameters for tools from the query.

        Args:
            query: The user's query string.
            required_tools: List of required tools.

        Returns:
            Dictionary of tool parameters.
        """
        params: dict[str, Any] = {}

        # Extract numeric values with units
        numeric_pattern = re.compile(
            r"(\d+\.?\d*)\s*(kg|g|mol|L|liters?|cm|m|mm|%|wt%|atom%)",
            re.IGNORECASE,
        )
        matches = numeric_pattern.findall(query)
        if matches:
            params["numeric_values"] = [
                {"value": float(m[0]), "unit": m[1]} for m in matches
            ]

        # Extract material mentions
        material_pattern = re.compile(
            r"\b(u-?235|u-?238|pu-?239|pu-?240|u-?233|uranium|plutonium)\b",
            re.IGNORECASE,
        )
        material_matches = material_pattern.findall(query)
        if material_matches:
            params["materials"] = list(set(m.lower().replace("-", "") for m in material_matches))

        # Extract geometry mentions
        geometry_pattern = re.compile(
            r"\b(sphere|cylinder|slab|cube|array|lattice)\b",
            re.IGNORECASE,
        )
        geometry_matches = geometry_pattern.findall(query)
        if geometry_matches:
            params["geometries"] = list(set(g.lower() for g in geometry_matches))

        # Extract enrichment values
        enrichment_pattern = re.compile(
            r"(\d+\.?\d*)\s*(%|percent|wt%|atom%)\s*(enriched?|enrichment)?",
            re.IGNORECASE,
        )
        enrichment_matches = enrichment_pattern.findall(query)
        if enrichment_matches:
            params["enrichment"] = float(enrichment_matches[0][0])

        return params

    def _estimate_reasoning_steps(
        self, complexity: ComplexityLevel, query_type: QueryType
    ) -> int:
        """Estimate the number of reasoning steps needed.

        Args:
            complexity: Query complexity level.
            query_type: Query type.

        Returns:
            Estimated number of reasoning steps.
        """
        base_steps = {
            ComplexityLevel.SIMPLE: 1,
            ComplexityLevel.MODERATE: 2,
            ComplexityLevel.COMPLEX: 4,
        }

        type_modifier = {
            QueryType.KNOWLEDGE: 0,
            QueryType.CALCULATION: 1,
            QueryType.DOCUMENT: 2,
            QueryType.TRAINING: 1,
            QueryType.ANALYSIS: 2,
        }

        return base_steps[complexity] + type_modifier[query_type]

    def _needs_rag(self, query_type: QueryType, required_tools: list[str]) -> bool:
        """Determine if RAG context is needed.

        Args:
            query_type: Query type.
            required_tools: List of required tools.

        Returns:
            True if RAG context should be retrieved.
        """
        # Almost all queries benefit from RAG context
        # Only pure unit conversions might skip RAG
        if query_type == QueryType.CALCULATION and required_tools == ["unit_converter"]:
            return False
        return True

    def _needs_verification(
        self, query_type: QueryType, complexity: ComplexityLevel
    ) -> bool:
        """Determine if response verification is needed.

        Args:
            query_type: Query type.
            complexity: Query complexity level.

        Returns:
            True if verification should be performed.
        """
        # Complex queries and calculations need verification
        if complexity == ComplexityLevel.COMPLEX:
            return True
        if query_type in (QueryType.CALCULATION, QueryType.ANALYSIS):
            return True
        return False

    def _get_routing_reasons(
        self,
        query: str,
        query_type: QueryType,
        required_tools: list[str],
    ) -> list[str]:
        """Generate explanations for routing decision.

        Args:
            query: The user's query string.
            query_type: Classified query type.
            required_tools: List of required tools.

        Returns:
            List of reason strings.
        """
        reasons = []

        type_explanations = {
            QueryType.KNOWLEDGE: "Query requests factual information",
            QueryType.CALCULATION: "Query involves numerical computation",
            QueryType.DOCUMENT: "Query requests document generation",
            QueryType.TRAINING: "Query requests educational content",
            QueryType.ANALYSIS: "Query requires multi-step analysis",
        }
        reasons.append(type_explanations[query_type])

        if required_tools:
            tool_names = ", ".join(required_tools)
            reasons.append(f"Tools identified: {tool_names}")

        return reasons
