"""Tool Runner for SAGE Engine.

Orchestrates execution of NCS tools (k-eff estimator, unit converter,
SCALE interface, standards lookup) with unified interface and error handling.
"""

import asyncio
import logging
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

from sage.engine.config import ToolConfig
from sage.tools import (
    GeometryVisualizer,
    KEffEstimator,
    SCALEInterface,
    StandardsLookup,
    UnitConverter,
)

logger = logging.getLogger(__name__)


class ToolStatus(str, Enum):
    """Status of a tool execution."""

    SUCCESS = "success"
    ERROR = "error"
    TIMEOUT = "timeout"
    UNAVAILABLE = "unavailable"


@dataclass
class ToolResult:
    """Result from a tool execution."""

    tool_name: str
    status: ToolStatus
    result: Any = None
    error: str | None = None
    execution_time_ms: float = 0.0
    warnings: list[str] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def success(self) -> bool:
        """Check if tool execution was successful."""
        return self.status == ToolStatus.SUCCESS


@dataclass
class ToolExecutionContext:
    """Context for tool execution."""

    user_id: str = "anonymous"
    request_id: str | None = None
    allow_dangerous: bool = False
    timeout_ms: float = 30000.0


class ToolRunner:
    """Orchestrates execution of NCS tools.

    Provides a unified interface for running k-eff estimator, unit converter,
    SCALE interface, and standards lookup tools with consistent error handling,
    timeouts, and result formatting.

    Example:
        runner = ToolRunner()

        # Run k-eff estimation
        result = await runner.run_tool(
            "keff_estimator",
            {"material": "U-235", "mass_kg": 10.0, "form": "metal"}
        )

        if result.success:
            print(f"k-eff = {result.result.k_eff_estimate}")

        # Run multiple tools
        results = await runner.run_tools([
            ("keff_estimator", {"material": "U-235", "mass_kg": 10.0}),
            ("unit_converter", {"value": 10.0, "from_unit": "kg", "to_unit": "g"}),
        ])
    """

    # Mapping of tool names to their classes and methods
    TOOL_REGISTRY: dict[str, dict[str, Any]] = {
        "keff_estimator": {
            "class": KEffEstimator,
            "method": "estimate",
            "description": "Estimate k-eff for fissile material configurations",
        },
        "unit_converter": {
            "class": UnitConverter,
            "method": "convert",
            "description": "Convert between NCS-relevant units",
        },
        "scale_interface": {
            "class": SCALEInterface,
            "method": "build_input",
            "description": "Generate SCALE input files for criticality calculations",
        },
        "standards_lookup": {
            "class": StandardsLookup,
            "method": "lookup",
            "description": "Look up NCS standards and subcritical limits",
        },
        "geometry_visualizer": {
            "class": GeometryVisualizer,
            "method": "visualize",
            "description": "Generate SVG visualizations of fissile material geometry configurations",
        },
    }

    def __init__(self, config: ToolConfig | None = None):
        """Initialize the tool runner.

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

        # Initialize tools lazily
        self._tools: dict[str, Any] = {}
        self._initialized = False

    def _get_tool(self, tool_name: str) -> Any:
        """Get or create a tool instance.

        Args:
            tool_name: Name of the tool.

        Returns:
            Tool instance.

        Raises:
            ValueError: If tool is not found or disabled.
        """
        if tool_name not in self.TOOL_REGISTRY:
            raise ValueError(f"Unknown tool: {tool_name}")

        # Check if tool is enabled
        enabled_flag = f"enable_{tool_name}"
        if hasattr(self.config, enabled_flag) and not getattr(self.config, enabled_flag):
            raise ValueError(f"Tool is disabled: {tool_name}")

        # Create instance if not cached
        if tool_name not in self._tools:
            tool_class = self.TOOL_REGISTRY[tool_name]["class"]
            self._tools[tool_name] = tool_class()
            logger.debug(f"Initialized tool: {tool_name}")

        return self._tools[tool_name]

    async def run_tool(
        self,
        tool_name: str,
        params: dict[str, Any],
        context: ToolExecutionContext | None = None,
    ) -> ToolResult:
        """Run a single tool with the given parameters.

        Args:
            tool_name: Name of the tool to run.
            params: Parameters to pass to the tool.
            context: Execution context with user info and settings.

        Returns:
            ToolResult with execution status and output.
        """
        import time

        context = context or ToolExecutionContext()
        start_time = time.time()

        try:
            # Get tool instance
            tool = self._get_tool(tool_name)

            # Get the method to call
            method_name = self.TOOL_REGISTRY[tool_name]["method"]

            # Handle special method routing based on params
            method_name = self._select_method(tool_name, method_name, params)
            method = getattr(tool, method_name)

            # Execute with timeout
            timeout_sec = min(
                context.timeout_ms / 1000,
                self.config.max_tool_execution_time_seconds,
            )

            # Run tool (synchronously, wrapped in asyncio)
            result = await asyncio.wait_for(
                asyncio.get_event_loop().run_in_executor(
                    None, lambda: self._execute_tool(method, params)
                ),
                timeout=timeout_sec,
            )

            execution_time = (time.time() - start_time) * 1000

            # Extract warnings if available
            warnings = []
            if hasattr(result, "warnings"):
                warnings = result.warnings

            return ToolResult(
                tool_name=tool_name,
                status=ToolStatus.SUCCESS,
                result=result,
                execution_time_ms=execution_time,
                warnings=warnings,
                metadata={"method": method_name, "params": params},
            )

        except asyncio.TimeoutError:
            execution_time = (time.time() - start_time) * 1000
            logger.warning(f"Tool {tool_name} timed out after {execution_time:.0f}ms")
            return ToolResult(
                tool_name=tool_name,
                status=ToolStatus.TIMEOUT,
                error=f"Tool execution timed out after {execution_time:.0f}ms",
                execution_time_ms=execution_time,
            )

        except ValueError as e:
            execution_time = (time.time() - start_time) * 1000
            if "disabled" in str(e).lower() or "unknown" in str(e).lower():
                return ToolResult(
                    tool_name=tool_name,
                    status=ToolStatus.UNAVAILABLE,
                    error=str(e),
                    execution_time_ms=execution_time,
                )
            return ToolResult(
                tool_name=tool_name,
                status=ToolStatus.ERROR,
                error=str(e),
                execution_time_ms=execution_time,
            )

        except Exception as e:
            execution_time = (time.time() - start_time) * 1000
            logger.error(f"Tool {tool_name} failed: {e}")
            return ToolResult(
                tool_name=tool_name,
                status=ToolStatus.ERROR,
                error=str(e),
                execution_time_ms=execution_time,
            )

    def _select_method(
        self, tool_name: str, default_method: str, params: dict[str, Any]
    ) -> str:
        """Select the appropriate method based on tool and params.

        Args:
            tool_name: Name of the tool.
            default_method: Default method name.
            params: Parameters provided.

        Returns:
            Method name to call.
        """
        # Unit converter has multiple methods
        if tool_name == "unit_converter":
            if "density" in params:
                return "mass_to_volume"
            if "concentration" in params and "volume" in params:
                return "concentration_to_mass"
            if "enrichment" in params and "from_format" in params:
                return "enrichment_format"
            return "convert"

        # K-eff estimator
        if tool_name == "keff_estimator":
            if params.get("validate_only"):
                return "validate_geometry"
            if params.get("check_margins"):
                return "check_margins"
            if params.get("get_limit"):
                return "get_limit"
            return "estimate"

        # SCALE interface
        if tool_name == "scale_interface":
            if params.get("validate_only"):
                return "validate_input"
            if params.get("parse_output"):
                return "parse_output"
            if params.get("generate"):
                return "generate_input"
            return "build_input"

        # Standards lookup
        if tool_name == "standards_lookup":
            if params.get("get_limits"):
                return "get_limits"
            if params.get("check_compliance"):
                return "check_compliance"
            if params.get("list_standards"):
                return "list_standards"
            return "lookup"

        # Geometry visualizer
        if tool_name == "geometry_visualizer":
            return "visualize"

        return default_method

    def _execute_tool(self, method: Any, params: dict[str, Any]) -> Any:
        """Execute a tool method with parameters.

        Args:
            method: Method to call.
            params: Parameters to pass.

        Returns:
            Method result.
        """
        # Filter out control params
        control_params = {
            "validate_only",
            "check_margins",
            "get_limit",
            "parse_output",
            "generate",
            "get_limits",
            "check_compliance",
            "list_standards",
        }
        filtered_params = {k: v for k, v in params.items() if k not in control_params}

        return method(**filtered_params)

    async def run_tools(
        self,
        tool_calls: list[tuple[str, dict[str, Any]]],
        context: ToolExecutionContext | None = None,
    ) -> list[ToolResult]:
        """Run multiple tools, potentially in parallel.

        Args:
            tool_calls: List of (tool_name, params) tuples.
            context: Execution context.

        Returns:
            List of ToolResult objects in same order as input.
        """
        context = context or ToolExecutionContext()

        # Limit concurrent executions
        semaphore = asyncio.Semaphore(self.config.max_concurrent_tools)

        async def run_with_semaphore(
            tool_name: str, params: dict[str, Any]
        ) -> ToolResult:
            async with semaphore:
                return await self.run_tool(tool_name, params, context)

        # Run all tools concurrently
        tasks = [run_with_semaphore(name, params) for name, params in tool_calls]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # Convert exceptions to ToolResults
        final_results = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                tool_name = tool_calls[i][0]
                final_results.append(
                    ToolResult(
                        tool_name=tool_name,
                        status=ToolStatus.ERROR,
                        error=str(result),
                    )
                )
            else:
                final_results.append(result)

        return final_results

    def select_tools(self, query: str, tool_params: dict[str, Any]) -> list[str]:
        """Select appropriate tools based on query analysis.

        Args:
            query: The user's query string.
            tool_params: Extracted parameters from query routing.

        Returns:
            List of recommended tool names.
        """
        selected = []

        # Check for materials/mass -> k-eff estimator
        if tool_params.get("materials") or tool_params.get("numeric_values"):
            has_mass = any(
                v.get("unit", "").lower() in ("kg", "g", "lb")
                for v in tool_params.get("numeric_values", [])
            )
            if has_mass or tool_params.get("materials"):
                if self.config.enable_keff_estimator:
                    selected.append("keff_estimator")

        # Check for unit conversion needs
        numeric_values = tool_params.get("numeric_values", [])
        if len(numeric_values) >= 1:
            if self.config.enable_unit_converter:
                selected.append("unit_converter")

        # Check for standards references
        query_lower = query.lower()
        if any(
            term in query_lower
            for term in ["ansi", "ans-8", "standard", "limit", "10 cfr"]
        ):
            if self.config.enable_standards_lookup:
                selected.append("standards_lookup")

        # Check for SCALE-related queries
        if any(
            term in query_lower
            for term in ["scale", "keno", "input file", "monte carlo"]
        ):
            if self.config.enable_scale_interface:
                selected.append("scale_interface")

        # Check for geometry visualization queries
        if any(
            term in query_lower
            for term in ["visualize", "diagram", "draw", "show geometry", "cross-section", "svg"]
        ):
            if self.config.enable_geometry_visualizer:
                selected.append("geometry_visualizer")

        return selected

    def get_available_tools(self) -> dict[str, str]:
        """Get dictionary of available tools and descriptions.

        Returns:
            Dict mapping tool names to descriptions.
        """
        available = {}
        for name, info in self.TOOL_REGISTRY.items():
            enabled_flag = f"enable_{name}"
            if not hasattr(self.config, enabled_flag) or getattr(
                self.config, enabled_flag
            ):
                available[name] = info["description"]
        return available

    def format_tool_results(self, results: list[ToolResult]) -> str:
        """Format tool results for inclusion in response.

        Args:
            results: List of tool results.

        Returns:
            Formatted string summarizing tool outputs.
        """
        if not results:
            return ""

        lines = ["**Tool Results:**", ""]

        for result in results:
            if result.success:
                lines.append(f"- **{result.tool_name}**: {self._format_result(result)}")
            else:
                lines.append(
                    f"- **{result.tool_name}**: Error - {result.error}"
                )

            if result.warnings:
                for warning in result.warnings:
                    lines.append(f"  - ⚠️ {warning}")

        return "\n".join(lines)

    def _format_result(self, result: ToolResult) -> str:
        """Format a single tool result for display.

        Args:
            result: Tool result to format.

        Returns:
            Formatted string.
        """
        data = result.result

        # Handle k-eff estimator result
        if result.tool_name == "keff_estimator" and hasattr(data, "k_eff_estimate"):
            status = "subcritical" if data.is_subcritical else "potentially critical"
            return (
                f"k-eff ≈ {data.k_eff_estimate:.3f} ± {data.uncertainty:.3f} ({status})"
            )

        # Handle unit converter result
        if result.tool_name == "unit_converter" and hasattr(data, "value"):
            return f"{data.value:.4g} {data.to_unit}"

        # Handle standards lookup result
        if result.tool_name == "standards_lookup":
            if hasattr(data, "standard_id"):
                return f"{data.standard_id}: {data.title}"
            if isinstance(data, list):
                return f"Found {len(data)} matching limits"
            if isinstance(data, tuple):
                return f"{'Compliant' if data[0] else 'Non-compliant'}: {data[1]}"

        # Handle SCALE interface result
        if result.tool_name == "scale_interface":
            if hasattr(data, "title"):
                return f"Generated input: {data.title}"
            if hasattr(data, "k_eff"):
                return f"k-eff = {data.k_eff:.5f} ± {data.k_eff_sigma:.5f}"

        # Handle geometry visualizer result
        if result.tool_name == "geometry_visualizer":
            if hasattr(data, "geometry_type"):
                dims = ", ".join(f"{k}={v:.1f}" for k, v in data.dimensions.items())
                return f"Generated {data.geometry_type} visualization ({dims} cm)"

        # Fallback
        return str(data)[:100]
