"""Instance manager - manages the lifecycle of MCP server instances."""

from __future__ import annotations

import asyncio
import contextlib
import logging
import time
from datetime import datetime
from typing import Any

import psutil

from mcp_bridge.config.settings import get_settings
from mcp_bridge.credentials.store import CredentialStore, create_default_store
from mcp_bridge.instances.connector import MCPConnection
from mcp_bridge.metrics.collector import MetricsCollector
from mcp_bridge.registry.catalog import ServerCatalog
from mcp_bridge.registry.models import InstanceState, ServerInstance, ToolCallResult, ToolInfo

logger = logging.getLogger(__name__)


class InstanceManager:
    """
    Manages the lifecycle of MCP server instances.

    Responsibilities:
    - Start/stop server instances on demand
    - Track instance state and ownership
    - Enforce resource limits (per-server, per-agent, system-wide)
    - Health monitoring and auto-restart
    - Idle timeout auto-shutdown
    - Route tool calls to the correct instance
    """

    def __init__(
        self,
        catalog: ServerCatalog,
        metrics: MetricsCollector,
        credential_store: CredentialStore | None = None,
    ):
        self._catalog = catalog
        self._metrics = metrics
        self._credentials = credential_store or create_default_store()

        # Active instances: instance_id -> (ServerInstance, MCPConnection)
        self._instances: dict[str, tuple[ServerInstance, MCPConnection]] = {}

        # Background tasks
        self._health_check_task: asyncio.Task | None = None
        self._idle_check_task: asyncio.Task | None = None
        self._running = False

    @property
    def active_instance_count(self) -> int:
        return sum(1 for inst, _ in self._instances.values() if inst.state == InstanceState.ACTIVE)

    async def start(self) -> None:
        """Start the instance manager background tasks."""
        self._running = True
        settings = get_settings()

        self._health_check_task = asyncio.create_task(
            self._health_check_loop(settings.instances.health_check_interval_seconds)
        )
        self._idle_check_task = asyncio.create_task(self._idle_check_loop())
        logger.info("Instance manager started")

    async def stop(self) -> None:
        """Stop all instances and background tasks."""
        self._running = False

        # Cancel background tasks
        for task in [self._health_check_task, self._idle_check_task]:
            if task and not task.done():
                task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await task

        # Shutdown all active instances
        instance_ids = list(self._instances.keys())
        for instance_id in instance_ids:
            try:
                await self.deactivate_instance(instance_id)
            except Exception as e:
                logger.error("Error shutting down instance %s: %s", instance_id, e)

        logger.info("Instance manager stopped")

    # --- Resource Checks ---

    def _check_system_resources(self) -> tuple[bool, str]:
        """Check if system resources are available for a new instance."""
        settings = get_settings()
        if not settings.system.check_resources_before_start:
            return True, "Resource checks disabled"

        cpu_percent = psutil.cpu_percent(interval=0.5)
        memory = psutil.virtual_memory()

        if cpu_percent > settings.system.max_cpu_percent:
            return False, f"CPU usage too high: {cpu_percent:.1f}% (max: {settings.system.max_cpu_percent}%)"

        if memory.percent > settings.system.max_memory_percent:
            return False, f"Memory usage too high: {memory.percent:.1f}% (max: {settings.system.max_memory_percent}%)"

        return True, "OK"

    def _check_instance_limits(self, server_id: str, agent_session_id: str) -> tuple[bool, str]:
        """Check if instance limits allow creating a new instance."""
        settings = get_settings()

        # Total instances
        total_active = self.active_instance_count
        if total_active >= settings.instances.max_total_instances:
            return False, f"Total instance limit reached: {total_active}/{settings.instances.max_total_instances}"

        # Per-server limit
        server_count = sum(
            1
            for inst, _ in self._instances.values()
            if inst.server_id == server_id and inst.state == InstanceState.ACTIVE
        )
        max_per_server = settings.instances.max_instances_per_server
        if server_count >= max_per_server:
            return False, f"Server '{server_id}' instance limit: {server_count}/{max_per_server}"

        # Per-agent limit
        agent_count = sum(
            1
            for inst, _ in self._instances.values()
            if inst.agent_session_id == agent_session_id and inst.state == InstanceState.ACTIVE
        )
        max_per_agent = settings.instances.max_instances_per_agent
        if agent_count >= max_per_agent:
            return False, f"Agent session instance limit: {agent_count}/{max_per_agent}"

        return True, "OK"

    # --- Instance Lifecycle ---

    async def configure_instance(
        self,
        server_id: str,
        agent_session_id: str,
        config: dict[str, Any] | None = None,
    ) -> ServerInstance:
        """
        Configure a new server instance without starting it.

        Returns the instance with a unique ID that can be used to activate it.
        """
        # Get server entry
        server_entry = await self._catalog.get_server(server_id)
        if server_entry is None:
            raise ValueError(f"Unknown server: {server_id}")

        if not server_entry.enabled:
            raise ValueError(f"Server '{server_id}' is disabled")

        # Check limits
        ok, msg = self._check_instance_limits(server_id, agent_session_id)
        if not ok:
            raise RuntimeError(msg)

        # Validate config against configurable options
        final_config = config or {}
        for option in server_entry.configurable_options:
            if option.name not in final_config and option.default is not None:
                final_config[option.name] = option.default
            if option.required and option.name not in final_config:
                raise ValueError(f"Required config option '{option.name}' not provided for server '{server_id}'")

        # Create instance
        instance = ServerInstance(
            server_id=server_id,
            agent_session_id=agent_session_id,
            state=InstanceState.CONFIGURED,
            config=final_config,
        )

        # Resolve credentials for required env vars
        resolved_env = dict(server_entry.env)
        for env_var in server_entry.required_env:
            value = await self._credentials.get(env_var)
            if value:
                resolved_env[env_var] = value

        # Store with a placeholder connection (not yet connected)
        connection = MCPConnection(server_entry, final_config)
        self._instances[instance.instance_id] = (instance, connection)

        logger.info("Configured instance %s for server '%s'", instance.instance_id, server_id)
        return instance

    async def activate_instance(self, instance_id: str) -> ServerInstance:
        """
        Activate (start) a configured instance.

        Connects to the MCP server, discovers tools, and makes them available.
        """
        if instance_id not in self._instances:
            raise ValueError(f"Unknown instance: {instance_id}")

        instance, connection = self._instances[instance_id]

        if instance.state == InstanceState.ACTIVE:
            return instance  # Already active

        if instance.state not in (InstanceState.CONFIGURED, InstanceState.STOPPED, InstanceState.ERROR):
            raise RuntimeError(f"Cannot activate instance in state '{instance.state}'")

        # Check system resources
        ok, msg = self._check_system_resources()
        if not ok:
            raise RuntimeError(msg)

        instance.state = InstanceState.STARTING
        instance.started_at = datetime.now()

        try:
            # Connect and discover tools
            tools = await asyncio.wait_for(
                connection.connect(),
                timeout=get_settings().instances.startup_timeout_seconds,
            )

            # Apply tool filtering
            server_entry = await self._catalog.get_server(instance.server_id)
            if server_entry and server_entry.tool_filter:
                tools = self._apply_tool_filter(tools, server_entry.tool_filter)

            instance.tools = tools
            instance.state = InstanceState.ACTIVE
            instance.last_activity_at = datetime.now()
            instance.pid = connection.pid

            if connection.port:
                instance.port = connection.port

            # Record metric
            await self._metrics.record_instance_event(
                instance_id=instance_id,
                server_id=instance.server_id,
                agent_session_id=instance.agent_session_id,
                event_type="activated",
                details={"tools_count": len(tools), "pid": instance.pid},
            )

            logger.info(
                "Activated instance %s (%s) with %d tools",
                instance_id,
                instance.server_id,
                len(tools),
            )
            return instance

        except TimeoutError:
            instance.state = InstanceState.ERROR
            raise RuntimeError(f"Server '{instance.server_id}' startup timed out") from None
        except Exception as e:
            instance.state = InstanceState.ERROR
            raise RuntimeError(f"Failed to activate instance: {e}") from e

    def _apply_tool_filter(self, tools: list[ToolInfo], tool_filter) -> list[ToolInfo]:
        """Apply whitelist/blacklist filtering to discovered tools."""
        if tool_filter.enabled_tools is not None:
            tools = [t for t in tools if t.name in tool_filter.enabled_tools]
        if tool_filter.disabled_tools is not None:
            tools = [t for t in tools if t.name not in tool_filter.disabled_tools]
        return tools

    async def deactivate_instance(self, instance_id: str) -> bool:
        """Deactivate (stop) a running instance."""
        if instance_id not in self._instances:
            return False

        instance, connection = self._instances[instance_id]
        instance.state = InstanceState.SHUTTING_DOWN

        try:
            await connection.disconnect()
        except Exception as e:
            logger.error("Error disconnecting instance %s: %s", instance_id, e)

        instance.state = InstanceState.STOPPED
        instance.tools = []

        # Record metric
        await self._metrics.record_instance_event(
            instance_id=instance_id,
            server_id=instance.server_id,
            agent_session_id=instance.agent_session_id,
            event_type="deactivated",
            details={"total_calls": instance.total_tool_calls},
        )

        # Remove from tracking
        del self._instances[instance_id]

        logger.info("Deactivated instance %s (%s)", instance_id, instance.server_id)
        return True

    # --- Tool Execution ---

    async def call_tool(
        self,
        instance_id: str,
        tool_name: str,
        arguments: dict[str, Any],
        agent_session_id: str,
    ) -> ToolCallResult:
        """Execute a tool call on a specific instance."""
        if instance_id not in self._instances:
            return ToolCallResult(
                instance_id=instance_id,
                tool_name=tool_name,
                success=False,
                error=f"Unknown instance: {instance_id}",
            )

        instance, connection = self._instances[instance_id]

        if instance.state != InstanceState.ACTIVE:
            return ToolCallResult(
                instance_id=instance_id,
                tool_name=tool_name,
                success=False,
                error=f"Instance not active (state: {instance.state})",
            )

        # Verify ownership
        if instance.agent_session_id != agent_session_id:
            return ToolCallResult(
                instance_id=instance_id,
                tool_name=tool_name,
                success=False,
                error="Instance belongs to a different agent session",
            )

        # Verify tool exists
        tool_names = [t.name for t in instance.tools]
        if tool_name not in tool_names:
            return ToolCallResult(
                instance_id=instance_id,
                tool_name=tool_name,
                success=False,
                error=f"Tool '{tool_name}' not found. Available: {tool_names}",
            )

        # Execute
        start_time = time.monotonic()
        try:
            result = await connection.call_tool(tool_name, arguments)
            latency_ms = (time.monotonic() - start_time) * 1000

            instance.total_tool_calls += 1
            instance.last_activity_at = datetime.now()

            # Record metric
            await self._metrics.record_tool_call(
                agent_session_id=agent_session_id,
                instance_id=instance_id,
                server_id=instance.server_id,
                tool_name=tool_name,
                success=True,
                latency_ms=latency_ms,
                input_size_bytes=len(str(arguments)),
                output_size_bytes=len(str(result)),
            )

            return ToolCallResult(
                instance_id=instance_id,
                tool_name=tool_name,
                success=True,
                result=result,
                latency_ms=round(latency_ms, 2),
            )

        except Exception as e:
            latency_ms = (time.monotonic() - start_time) * 1000
            instance.total_errors += 1
            instance.last_activity_at = datetime.now()

            await self._metrics.record_tool_call(
                agent_session_id=agent_session_id,
                instance_id=instance_id,
                server_id=instance.server_id,
                tool_name=tool_name,
                success=False,
                latency_ms=latency_ms,
                error_type=type(e).__name__,
                error_message=str(e),
            )

            return ToolCallResult(
                instance_id=instance_id,
                tool_name=tool_name,
                success=False,
                error=str(e),
                latency_ms=round(latency_ms, 2),
            )

    async def parallel_call(
        self,
        calls: list[dict[str, Any]],
        agent_session_id: str,
    ) -> list[ToolCallResult]:
        """Execute multiple tool calls in parallel across different instances."""
        tasks = []
        for call in calls:
            task = self.call_tool(
                instance_id=call["instance_id"],
                tool_name=call["tool_name"],
                arguments=call.get("arguments", {}),
                agent_session_id=agent_session_id,
            )
            tasks.append(task)

        results = await asyncio.gather(*tasks, return_exceptions=True)

        # Convert exceptions to ToolCallResults
        final_results = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                final_results.append(
                    ToolCallResult(
                        instance_id=calls[i]["instance_id"],
                        tool_name=calls[i]["tool_name"],
                        success=False,
                        error=str(result),
                    )
                )
            else:
                final_results.append(result)

        return final_results

    # --- Query ---

    def get_instance(self, instance_id: str) -> ServerInstance | None:
        """Get instance info by ID."""
        if instance_id in self._instances:
            return self._instances[instance_id][0]
        return None

    def get_instances_for_agent(self, agent_session_id: str) -> list[ServerInstance]:
        """Get all instances owned by an agent session."""
        return [inst for inst, _ in self._instances.values() if inst.agent_session_id == agent_session_id]

    def get_all_instances(self) -> list[ServerInstance]:
        """Get all active instances."""
        return [inst for inst, _ in self._instances.values()]

    def get_instance_tools(self, instance_id: str) -> list[ToolInfo]:
        """Get the tools available on an instance."""
        if instance_id in self._instances:
            return self._instances[instance_id][0].tools
        return []

    # --- Background Tasks ---

    async def _health_check_loop(self, interval_seconds: int) -> None:
        """Periodically check health of active instances."""
        while self._running:
            try:
                await asyncio.sleep(interval_seconds)
                for instance_id, (instance, connection) in list(self._instances.items()):
                    if instance.state != InstanceState.ACTIVE:
                        continue

                    healthy = await connection.health_check()
                    if not healthy:
                        logger.warning("Instance %s (%s) failed health check", instance_id, instance.server_id)
                        instance.state = InstanceState.ERROR

                        # Auto-restart if configured
                        settings = get_settings()
                        if instance.total_errors < settings.instances.max_restart_attempts:
                            logger.info("Attempting to restart instance %s", instance_id)
                            try:
                                await connection.disconnect()
                                await connection.connect()
                                instance.state = InstanceState.ACTIVE
                                logger.info("Successfully restarted instance %s", instance_id)
                            except Exception as e:
                                logger.error("Failed to restart instance %s: %s", instance_id, e)

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error("Error in health check loop: %s", e)

    async def _idle_check_loop(self) -> None:
        """Check for idle instances and shut them down."""
        while self._running:
            try:
                await asyncio.sleep(30)  # Check every 30 seconds
                settings = get_settings()
                now = datetime.now()

                for instance_id, (instance, _connection) in list(self._instances.items()):
                    if instance.state != InstanceState.ACTIVE:
                        continue

                    # Get idle timeout
                    server_entry = await self._catalog.get_server(instance.server_id)
                    idle_timeout = (
                        server_entry.auto_shutdown_idle_seconds
                        if server_entry and server_entry.auto_shutdown_idle_seconds
                        else settings.instances.idle_timeout_seconds
                    )

                    last_activity = instance.last_activity_at or instance.started_at or instance.created_at
                    idle_seconds = (now - last_activity).total_seconds()

                    if idle_seconds > idle_timeout:
                        logger.info(
                            "Instance %s (%s) idle for %.0fs, auto-shutting down",
                            instance_id,
                            instance.server_id,
                            idle_seconds,
                        )
                        await self.deactivate_instance(instance_id)

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error("Error in idle check loop: %s", e)
