"""MCP integration for HLA-Compass modules.

Enables any module to be exposed as an MCP (Model Context Protocol) tool,
allowing AI clients to discover and invoke modules programmatically.

Usage::

    from hla_compass.mcp import MCPModuleServer

    class MyModule(Module):
        ...

    server = MCPModuleServer(MyModule, manifest_path="manifest.json")
    server.run()  # starts stdio MCP server

Generating a tool definition without starting a server::

    from hla_compass.mcp import module_to_mcp_tool

    tool = module_to_mcp_tool(manifest)
    print(json.dumps(tool, indent=2))
"""

from __future__ import annotations

import json
import logging
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Optional, Type

logger = logging.getLogger(__name__)


def module_to_mcp_tool(manifest: Dict[str, Any]) -> Dict[str, Any]:
    """Generate an MCP tool definition from a module manifest.

    Args:
        manifest: Parsed ``manifest.json`` content.

    Returns:
        An MCP-compatible tool definition dict with ``name``,
        ``description``, and ``inputSchema``.
    """
    inputs = manifest.get("inputs", {})

    # Convert manifest inputs to JSON Schema (MCP inputSchema)
    if isinstance(inputs, dict):
        if inputs.get("type") == "object" and "properties" in inputs:
            input_schema = inputs
        else:
            input_schema = _flat_to_json_schema(inputs)
    else:
        input_schema = {"type": "object", "properties": {}}

    return {
        "name": manifest.get("name", "unknown-module"),
        "description": manifest.get("description", ""),
        "inputSchema": input_schema,
    }


def _flat_to_json_schema(flat_inputs: Dict[str, Any]) -> Dict[str, Any]:
    """Convert flat manifest input definitions to JSON Schema."""
    properties: Dict[str, Any] = {}
    required: list[str] = []

    for field_name, field_def in flat_inputs.items():
        if not isinstance(field_def, dict):
            continue
        prop: Dict[str, Any] = {}
        for key in ("type", "enum", "minimum", "maximum", "default", "description"):
            if key in field_def:
                prop[key] = field_def[key]
        if field_def.get("required"):
            required.append(field_name)
        properties[field_name] = prop

    schema: Dict[str, Any] = {"type": "object", "properties": properties}
    if required:
        schema["required"] = required
    return schema


def _make_mcp_context(manifest: Dict[str, Any]) -> Dict[str, Any]:
    """Create a minimal execution context for MCP tool invocations."""
    return {
        "run_id": f"mcp-{uuid.uuid4().hex[:8]}",
        "module_id": manifest.get("name", "mcp-module"),
        "module_version": manifest.get("version", "0.0.0"),
        "organization_id": "mcp-client",
        "user_id": "mcp-client",
        "environment": "local",
        "correlation_id": f"mcp-{uuid.uuid4().hex[:8]}",
        "requested_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        "roles": ["developer"],
        "mode": "interactive",
        "offline": True,
    }


class MCPModuleServer:
    """MCP server that wraps an HLA-Compass module.

    Each module is exposed as a single MCP tool whose ``inputSchema``
    is derived from the module manifest.

    Args:
        module_class: The :class:`~hla_compass.module.Module` subclass.
        manifest_path: Path to ``manifest.json``.
    """

    def __init__(
        self,
        module_class: Type,
        manifest_path: str = "manifest.json",
    ):
        self.module_class = module_class
        self.manifest_path = manifest_path
        self._manifest: Optional[Dict[str, Any]] = None

    def _load_manifest(self) -> Dict[str, Any]:
        if self._manifest is not None:
            return self._manifest
        path = Path(self.manifest_path)
        if path.exists():
            self._manifest = json.loads(path.read_text(encoding="utf-8"))
        else:
            self._manifest = {}
        return self._manifest

    def _get_name(self) -> str:
        return self._load_manifest().get("name", "hla-compass-module")

    def get_tool_definition(self) -> Dict[str, Any]:
        """Return the MCP tool definition for this module."""
        return module_to_mcp_tool(self._load_manifest())

    def run(self, transport: str = "stdio") -> None:
        """Start the MCP server.

        Args:
            transport: Transport type — ``"stdio"`` (default) for
                command-line usage with AI clients.

        Raises:
            ImportError: If the ``mcp`` package is not installed.
        """
        try:
            from mcp.server import Server
            from mcp.server.stdio import stdio_server
        except ImportError as err:
            raise ImportError(
                "MCP server requires the 'mcp' package. Install with: pip install 'hla-compass[mcp]'"
            ) from err

        manifest = self._load_manifest()
        tool_def = module_to_mcp_tool(manifest)
        server = Server(name=self._get_name())

        module_class = self.module_class
        manifest_path = self.manifest_path

        @server.tool(
            name=tool_def["name"],
            description=tool_def.get("description", ""),
            input_schema=tool_def["inputSchema"],
        )
        async def execute_module(arguments: dict) -> str:
            module = module_class(manifest_path=manifest_path)
            context = _make_mcp_context(manifest)
            result = module.run(arguments, context)
            return json.dumps(result, default=str)

        logger.info("Starting MCP server for module '%s' via %s", self._get_name(), transport)

        if transport == "stdio":
            import asyncio

            asyncio.run(stdio_server(server))
        else:
            raise ValueError(f"Unsupported transport: {transport!r}. Use 'stdio'.")
