"""Tool manager for MCP servers."""

import asyncio
import json
import os
import re
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any

import httpx

from cadecoder.core.logging import log

# Anthropic requires tool names to match this pattern
_VALID_TOOL_NAME_PATTERN = re.compile(r"^[a-zA-Z0-9_-]{1,128}$")


def sanitize_tool_name(name: str) -> str:
    """Sanitize a tool name to match API requirements.

    Anthropic requires tool names to match: ^[a-zA-Z0-9_-]{1,128}$

    Args:
        name: The original tool name

    Returns:
        A sanitized tool name that matches the required pattern
    """
    if _VALID_TOOL_NAME_PATTERN.match(name):
        return name

    # Replace invalid characters with underscores
    sanitized = re.sub(r"[^a-zA-Z0-9_-]", "_", name)

    # Remove consecutive underscores
    sanitized = re.sub(r"_+", "_", sanitized)

    # Strip leading/trailing underscores
    sanitized = sanitized.strip("_")

    # Ensure not empty and not too long
    if not sanitized:
        sanitized = "unnamed_tool"
    sanitized = sanitized[:128]

    return sanitized


# Base classes and exceptions


class ToolAuthorizationRequired(Exception):
    """Exception raised when tool execution requires user authorization."""

    def __init__(self, tool_name: str, authorization_url: str | None = None):
        self.tool_name = tool_name
        self.authorization_url = authorization_url
        super().__init__(f"Authorization required for tool: {tool_name}")


class ToolManager:
    """Base class for tool managers."""

    async def get_tools(self) -> list[dict[str, Any]]:
        """Get list of available tools."""
        raise NotImplementedError

    async def execute(self, name: str, inputs: dict[str, Any]) -> Any:
        """Execute a tool."""
        raise NotImplementedError

    async def close(self) -> None:
        """Close connections."""
        pass


# MCP Configuration


class MCPTransportType(str, Enum):
    """MCP transport types."""

    STDIO = "stdio"
    HTTP = "http"
    SSE = "sse"


class MCPAuthType(str, Enum):
    """MCP authentication types."""

    NONE = "none"
    BEARER = "bearer"
    API_KEY = "api_key"
    OAUTH = "oauth"


@dataclass
class MCPOAuthTokens:
    """OAuth tokens for MCP servers."""

    access_token: str
    token_type: str = "Bearer"
    refresh_token: str | None = None
    expires_at: datetime | None = None
    scope: str | None = None

    def is_expired(self) -> bool:
        """Check if access token is expired."""
        if not self.expires_at:
            return False
        return datetime.now() >= self.expires_at


@dataclass
class MCPServerConfig:
    """Configuration for an MCP server."""

    name: str
    url: str
    transport: MCPTransportType = MCPTransportType.HTTP
    command: str | None = None
    args: list[str] = field(default_factory=list)
    env: dict[str, str] = field(default_factory=dict)
    enabled: bool = True
    auth_type: MCPAuthType = MCPAuthType.NONE
    auth_value: str | None = None
    oauth_tokens: MCPOAuthTokens | None = None
    oauth_client_id: str | None = None
    oauth_client_secret: str | None = None
    oauth_authorization_server: str | None = None
    oauth_scopes: list[str] = field(default_factory=list)
    last_connected: datetime | None = None
    tool_count: int = 0

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "name": self.name,
            "url": self.url,
            "transport": self.transport.value,
            "auth_type": self.auth_type.value,
            "auth_value": self.auth_value,
            "enabled": self.enabled,
            "last_connected": self.last_connected.isoformat() if self.last_connected else None,
            "tool_count": self.tool_count,
            "command": self.command,
            "args": self.args if self.args else None,
            "env": self.env if self.env else None,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "MCPServerConfig":
        """Create from dictionary."""
        last_connected = None
        if data.get("last_connected"):
            last_connected = datetime.fromisoformat(data["last_connected"])

        return cls(
            name=data["name"],
            url=data["url"],
            transport=MCPTransportType(data.get("transport", "http")),
            auth_type=MCPAuthType(data.get("auth_type", "none")),
            auth_value=data.get("auth_value"),
            enabled=data.get("enabled", True),
            last_connected=last_connected,
            tool_count=data.get("tool_count", 0),
            command=data.get("command"),
            args=data.get("args") or [],
            env=data.get("env") or {},
        )


class MCPServerStore:
    """Persistent storage for MCP server configurations."""

    def __init__(self, config_dir: Path | None = None) -> None:
        from cadecoder.core.config import get_config

        if config_dir is None:
            config_dir = Path(get_config().app_dir)
        self.config_dir = config_dir
        self.config_file = config_dir / "mcp_servers.json"
        self._servers: dict[str, MCPServerConfig] = {}
        self._load()

    def _load(self) -> None:
        """Load servers from disk."""
        if not self.config_file.exists():
            return

        try:
            with open(self.config_file, encoding="utf-8") as f:
                data = json.load(f)

            for server_data in data.get("servers", []):
                server = MCPServerConfig.from_dict(server_data)
                self._servers[server.name] = server

            log.debug(f"Loaded {len(self._servers)} MCP server configs")
        except Exception as e:
            log.warning(f"Failed to load MCP server configs: {e}")

    def _save(self) -> None:
        """Save servers to disk."""
        self.config_dir.mkdir(parents=True, exist_ok=True)

        try:
            data = {
                "servers": [s.to_dict() for s in self._servers.values()],
                "updated_at": datetime.now().isoformat(),
            }
            with open(self.config_file, "w", encoding="utf-8") as f:
                json.dump(data, f, indent=2)
        except Exception as e:
            log.warning(f"Failed to save MCP server configs: {e}")

    def add(self, server: MCPServerConfig) -> None:
        """Add or update a server configuration."""
        self._servers[server.name] = server
        self._save()

    def remove(self, name: str) -> bool:
        """Remove a server configuration."""
        if name in self._servers:
            del self._servers[name]
            self._save()
            return True
        return False

    def get(self, name: str) -> MCPServerConfig | None:
        """Get a server configuration by name."""
        return self._servers.get(name)

    def list_all(self) -> list[MCPServerConfig]:
        """List all server configurations."""
        return list(self._servers.values())

    def list_enabled(self) -> list[MCPServerConfig]:
        """List all enabled server configurations."""
        return [s for s in self._servers.values() if s.enabled]

    def update_status(self, name: str, connected: bool, tool_count: int = 0) -> None:
        """Update server connection status."""
        if name in self._servers:
            if connected:
                self._servers[name].last_connected = datetime.now()
                self._servers[name].tool_count = tool_count
            self._save()


# Stdio MCP Connection


class StdioMCPConnection:
    """Manages stdio subprocess for MCP server."""

    def __init__(
        self, command: str, args: list[str] | None = None, env: dict[str, str] | None = None
    ):
        self.command = command
        self.args = args or []
        self.env = env
        self.process: asyncio.subprocess.Process | None = None
        self._lock = asyncio.Lock()

    async def start(self) -> None:
        """Start the stdio subprocess."""
        if self.process:
            return

        full_env = os.environ.copy()
        if self.env:
            full_env.update(self.env)

        log.debug(f"Starting stdio MCP server: {self.command} {' '.join(self.args)}")

        self.process = await asyncio.create_subprocess_exec(
            self.command,
            *self.args,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=full_env,
        )

        # Start stderr logger task
        asyncio.create_task(self._log_stderr())

    async def _log_stderr(self) -> None:
        """Log stderr output from subprocess."""
        if not self.process or not self.process.stderr:
            return

        try:
            async for line in self.process.stderr:
                log.debug(f"[stdio] {line.decode().rstrip()}")
        except Exception:
            pass

    async def send_request(self, request: dict[str, Any]) -> dict[str, Any]:
        """Send JSON-RPC request via stdio."""
        async with self._lock:
            if not self.process or not self.process.stdin or not self.process.stdout:
                raise Exception("Process not started")

            # Write request
            line = json.dumps(request) + "\n"
            self.process.stdin.write(line.encode())
            await self.process.stdin.drain()

            # Read response
            response_line = await self.process.stdout.readline()
            if not response_line:
                raise Exception("Process closed unexpectedly")

            return json.loads(response_line)

    async def close(self) -> None:
        """Close the subprocess."""
        if self.process:
            if self.process.stdin:
                self.process.stdin.close()
            try:
                await asyncio.wait_for(self.process.wait(), timeout=5.0)
            except TimeoutError:
                log.warning("Stdio process did not terminate gracefully, killing...")
                self.process.kill()
                await self.process.wait()
            self.process = None


# MCP Tool Manager


class MCPToolManager(ToolManager):
    """Manages tools from MCP servers over HTTP or stdio transport."""

    def __init__(self, server_config: MCPServerConfig):
        self.config = server_config
        self._tools_cache: list[dict[str, Any]] | None = None
        self._initialized = False

        # Transport-specific clients
        if server_config.transport == MCPTransportType.STDIO:
            self._stdio_connection = StdioMCPConnection(
                command=server_config.command or server_config.url,
                args=server_config.args,
                env=server_config.env,
            )
            self._client: httpx.AsyncClient | None = None
        else:
            self._stdio_connection: StdioMCPConnection | None = None
            self._client: httpx.AsyncClient | None = None

        self._session_id: str | None = None

    def _get_headers(self) -> dict[str, str]:
        """Get HTTP headers including authentication."""
        headers = {
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
        }

        if self._session_id:
            headers["Mcp-Session-Id"] = self._session_id

        if self.config.auth_type == MCPAuthType.BEARER and self.config.auth_value:
            headers["Authorization"] = f"Bearer {self.config.auth_value}"
        elif self.config.auth_type == MCPAuthType.API_KEY and self.config.auth_value:
            headers["X-API-Key"] = self.config.auth_value

        return headers

    async def _ensure_client(self) -> httpx.AsyncClient:
        """Ensure HTTP client is created."""
        if self._client is None:
            self._client = httpx.AsyncClient(timeout=httpx.Timeout(60.0, connect=10.0))
        return self._client

    async def _send_request(self, method: str, params: dict[str, Any] | None = None) -> Any:
        """Send a JSON-RPC request to the MCP server."""
        if self.config.transport == MCPTransportType.STDIO:
            return await self._send_stdio_request(method, params)
        else:
            return await self._send_http_request(method, params)

    async def _send_stdio_request(self, method: str, params: dict[str, Any] | None = None) -> Any:
        """Send a JSON-RPC request via stdio transport."""
        if not self._stdio_connection:
            raise Exception("Stdio connection not initialized")

        await self._stdio_connection.start()

        request_id = str(uuid.uuid4())
        payload: dict[str, Any] = {
            "jsonrpc": "2.0",
            "id": request_id,
            "method": method,
        }
        if params:
            payload["params"] = params

        try:
            response = await self._stdio_connection.send_request(payload)

            if "error" in response:
                raise Exception(f"MCP error: {response['error'].get('message', 'Unknown error')}")

            return response.get("result")

        except Exception as e:
            log.error(f"MCP stdio request failed for {self.config.name}: {e}")
            raise

    async def _send_http_request(self, method: str, params: dict[str, Any] | None = None) -> Any:
        """Send a JSON-RPC request via HTTP transport."""
        client = await self._ensure_client()

        request_id = str(uuid.uuid4())
        payload: dict[str, Any] = {
            "jsonrpc": "2.0",
            "id": request_id,
            "method": method,
        }
        if params:
            payload["params"] = params

        try:
            headers = self._get_headers()
            response = await client.post(self.config.url, json=payload, headers=headers)

            if "mcp-session-id" in response.headers:
                self._session_id = response.headers["mcp-session-id"]

            if response.status_code == 401:
                raise ToolAuthorizationRequired(f"MCP:{self.config.name}")

            response.raise_for_status()

            result = response.json()
            if "error" in result:
                raise Exception(f"MCP error: {result['error'].get('message', 'Unknown error')}")
            return result.get("result")

        except httpx.HTTPStatusError as e:
            raise Exception(f"MCP HTTP error: {e}")
        except ToolAuthorizationRequired:
            raise
        except Exception as e:
            log.error(f"MCP request failed for {self.config.name}: {e}")
            raise

    async def initialize(self) -> bool:
        """Initialize connection to MCP server."""
        if self._initialized:
            return True

        try:
            result = await self._send_request(
                "initialize",
                {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {"tools": {}},
                    "clientInfo": {"name": "cadecoder", "version": "1.0.0"},
                },
            )

            if result:
                # Send initialized notification (required by MCP protocol)
                await self._send_notification("notifications/initialized")
                self._initialized = True
                log.info(f"MCP server '{self.config.name}' initialized")
                return True

            return False

        except Exception as e:
            log.error(f"Failed to initialize MCP server '{self.config.name}': {e}")
            return False

    async def check_status(self) -> tuple[bool, str]:
        """Check connectivity status of the MCP server.

        Returns:
            Tuple of (connected: bool, status_message: str)
        """
        try:
            success = await self.initialize()
            if success:
                return True, "Connected"
            return False, "Failed to initialize"
        except Exception as e:
            return False, str(e)[:50]

    async def _send_notification(self, method: str, params: dict[str, Any] | None = None) -> None:
        """Send a JSON-RPC notification (no response expected)."""
        payload: dict[str, Any] = {
            "jsonrpc": "2.0",
            "method": method,
        }
        if params:
            payload["params"] = params

        if self.config.transport == MCPTransportType.STDIO:
            if not self._stdio_connection:
                raise Exception("Stdio connection not initialized")

            await self._stdio_connection.start()

            line = json.dumps(payload) + "\n"
            if self._stdio_connection.process and self._stdio_connection.process.stdin:
                self._stdio_connection.process.stdin.write(line.encode())
                await self._stdio_connection.process.stdin.drain()
        else:
            # HTTP transport - send notification as POST (no id = notification)
            client = await self._ensure_client()
            headers = self._get_headers()
            await client.post(self.config.url, json=payload, headers=headers)

    async def get_tools(self) -> list[dict[str, Any]]:
        """Get available tools from the MCP server."""
        if self._tools_cache is not None:
            return self._tools_cache

        if not self._initialized:
            success = await self.initialize()
            if not success:
                return []

        try:
            result = await self._send_request("tools/list")

            if not result or "tools" not in result:
                return []

            self._tools_cache = []
            for mcp_tool in result["tools"]:
                openai_tool = self._convert_mcp_to_openai(mcp_tool)
                self._tools_cache.append(openai_tool)

            log.info(f"Loaded {len(self._tools_cache)} tools from MCP server '{self.config.name}'")
            return self._tools_cache

        except Exception as e:
            log.error(f"Failed to list tools from MCP '{self.config.name}': {e}")
            return []

    def _convert_mcp_to_openai(self, mcp_tool: dict[str, Any]) -> dict[str, Any]:
        """Convert MCP tool schema to OpenAI function schema."""
        input_schema = mcp_tool.get("inputSchema", {"type": "object", "properties": {}})

        # Sanitize tool name to meet API requirements
        original_name = mcp_tool["name"]
        sanitized_name = sanitize_tool_name(original_name)

        if sanitized_name != original_name:
            log.debug(f"Sanitized tool name: '{original_name}' -> '{sanitized_name}'")

        return {
            "type": "function",
            "function": {
                "name": sanitized_name,
                "description": mcp_tool.get("description", ""),
                "parameters": input_schema,
            },
            # Store original name for execution routing
            "_original_name": original_name,
        }

    def _extract_text_from_content(self, content: list[dict[str, Any]]) -> str | None:
        """Extract text from MCP content array."""
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                return item.get("text")
        return None

    async def execute(self, name: str, inputs: dict[str, Any]) -> Any:
        """Execute a tool on the MCP server."""
        if not self._initialized:
            success = await self.initialize()
            if not success:
                raise Exception(f"MCP server '{self.config.name}' not initialized")

        try:
            result = await self._send_request(
                "tools/call",
                {"name": name, "arguments": inputs},
            )

            # Check if result is an error from arcade-mcp-server
            if isinstance(result, dict) and result.get("isError"):
                error_content = result.get("content", [])
                error_text = self._extract_text_from_content(error_content)
                if error_text:
                    raise Exception(error_text)
                raise Exception(f"Tool execution failed: {error_content}")

            # Check for successful result with content
            if isinstance(result, dict) and "content" in result:
                content_items = result["content"]
                text_parts = []
                for item in content_items:
                    if item.get("type") == "text":
                        text_parts.append(item.get("text", ""))
                return "\n".join(text_parts) if text_parts else result

            return result

        except Exception as e:
            log.error(f"MCP tool execution failed for '{name}': {e}")
            raise Exception(f"Failed to execute MCP tool '{name}': {e}")

    async def close(self) -> None:
        """Close connections."""
        if self._stdio_connection:
            await self._stdio_connection.close()
        if self._client:
            await self._client.aclose()
            self._client = None
        self._session_id = None
        self._initialized = False


# Composite Tool Manager


class CompositeToolManager(ToolManager):
    """Manages tools from all MCP servers (both local and external)."""

    def __init__(self, enable_mcp: bool = True, local_only: bool = False):
        """Initialize composite tool manager.

        Args:
            enable_mcp: Whether to enable external MCP servers
            local_only: If True, only load local tools (skip external MCP servers)
        """
        self._tool_to_manager: dict[str, MCPToolManager] = {}
        self._mcp_managers: list[MCPToolManager] = []
        self._enable_mcp = enable_mcp
        self._local_only = local_only
        self._tools_loaded = False
        # Maps sanitized tool names to original MCP names for execution
        self._sanitized_to_original: dict[str, str] = {}

        # Create local MCP server manager (always enabled)
        local_mcp_config = self._create_local_server_config()
        local_mcp_manager = MCPToolManager(local_mcp_config)
        self._mcp_managers.append(local_mcp_manager)

        # Load external MCP servers (skip if local_only mode)
        if local_only:
            log.info("Local-only mode: external MCP servers disabled")
        elif enable_mcp:
            self._load_external_mcp_servers()

    def _load_external_mcp_servers(self) -> None:
        """Load external MCP servers from the server store."""
        try:
            store = MCPServerStore()
            for server_config in store.list_enabled():
                # Skip if it's a duplicate of the local server
                if server_config.name == "local":
                    continue
                mcp_manager = MCPToolManager(server_config)
                self._mcp_managers.append(mcp_manager)
                log.debug(f"Loaded external MCP server: {server_config.name}")
        except Exception as e:
            log.warning(f"Failed to load external MCP servers: {e}")

    def _create_local_server_config(self) -> MCPServerConfig:
        """Create MCP config for local tools server via stdio.

        Uses the installed `cade-mcp-local` entry point and passes
        the current working directory via CADE_PROJECT_ROOT env var.
        """
        return MCPServerConfig(
            name="local",
            url="cade-mcp-local",
            transport=MCPTransportType.STDIO,
            command="cade-mcp-local",
            args=[],
            env={"CADE_PROJECT_ROOT": str(Path.cwd().resolve())},
            enabled=True,
            auth_type=MCPAuthType.NONE,
        )

    async def get_tools(self) -> list[dict[str, Any]]:
        """Get all available tools from all MCP servers."""
        all_tools = []

        for mcp_manager in self._mcp_managers:
            try:
                mcp_tools = await mcp_manager.get_tools()

                for tool in mcp_tools:
                    tool_name = tool.get("function", {}).get("name", "unknown")
                    original_name = tool.pop("_original_name", tool_name)

                    # Store mapping from sanitized name to manager and original name
                    self._tool_to_manager[tool_name] = mcp_manager
                    self._sanitized_to_original[tool_name] = original_name

                    all_tools.append(tool)

                log.info(
                    f"Loaded {len(mcp_tools)} tools from MCP server '{mcp_manager.config.name}'"
                )

            except Exception as e:
                log.error(f"Failed to load tools from MCP '{mcp_manager.config.name}': {e}")

        log.info(f"Total tools loaded: {len(all_tools)}")
        self._tools_loaded = True
        return all_tools

    async def execute(self, name: str, inputs: dict[str, Any]) -> Any:
        """Execute a tool by routing to the appropriate MCP manager."""
        if not self._tools_loaded:
            log.warning("Tools not loaded before execute() was called. Loading now...")
            await self.get_tools()

        manager = self._tool_to_manager.get(name)
        if not manager:
            available_tools = list(self._tool_to_manager.keys())[:10]
            raise Exception(
                f"Tool '{name}' not found. "
                f"Available tools (first 10): {available_tools}. "
                f"Total tools: {len(self._tool_to_manager)}"
            )

        # Use the original MCP tool name for execution
        original_name = self._sanitized_to_original.get(name, name)
        log.debug(
            f"Executing tool: {name} (original: {original_name}) via server '{manager.config.name}'"
        )
        return await manager.execute(original_name, inputs)

    def is_interactive_tool(self, name: str) -> bool:
        """Determine whether a tool is interactive."""
        return name in ("bash", "execute_command")

    def get_tool_source(self, name: str) -> str | None:
        """Get the source server name for a tool."""
        manager = self._tool_to_manager.get(name)
        return manager.config.name if manager else None

    def get_all_tool_info(self) -> list[dict[str, str]]:
        """Get info about all tools including their source servers.

        Returns:
            List of dicts with 'name' and 'server' keys for each tool.
        """
        return [
            {"name": name, "server": manager.config.name}
            for name, manager in self._tool_to_manager.items()
        ]

    async def close(self) -> None:
        """Close all MCP connections."""
        for mcp_manager in self._mcp_managers:
            await mcp_manager.close()
