"""Context MCP server for user context management."""

import logging
import os
import time
from typing import Any, Callable, Coroutine

from fastmcp import FastMCP

from .tools.retrieve import retrieve_user_context
from .tools.save import save_user_context

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def tool_wrapper(fn: Callable[..., Coroutine]) -> Callable[..., Coroutine]:
    """A decorator to wrap MCP tools with logging and error handling."""
    async def wrapper(*args, **kwargs):
        tool_name = fn.__name__
        customer_id = kwargs.get("customer_id", args[0] if args else "N/A")
        logger.info(
            "Invoking tool: %s",
            tool_name,
            extra={"tool_name": tool_name, "customer_id": customer_id}
        )
        start_time = time.perf_counter()
        try:
            result = await fn(*args, **kwargs)
            end_time = time.perf_counter()
            duration = (end_time - start_time) * 1000
            logger.info(
                "Tool %s completed successfully in %.2fms",
                tool_name,
                duration,
                extra={"tool_name": tool_name, "customer_id": customer_id, "duration_ms": duration}
            )
            return result
        except Exception as exc:  # pylint: disable=broad-exception-caught
            end_time = time.perf_counter()
            duration = (end_time - start_time) * 1000
            logger.error(
                "Tool %s failed after %.2fms: %s",
                tool_name,
                duration,
                exc,
                extra={
                    "tool_name": tool_name,
                    "customer_id": customer_id,
                    "duration_ms": duration,
                    "error": str(exc)
                },
                exc_info=True
            )
            # Return structured error response
            error_type = type(exc).__name__
            return {
                "success": False,
                "error_type": error_type,
                "error_message": str(exc),
                "user_message": f"An error occurred while processing your request: {exc!s}"
            }
    return wrapper

# Configuration

server_name = os.getenv("SERVER_NAME", "Context MCP")

mcp = FastMCP(
        name=server_name
    )


@mcp.tool()
async def save_context(customer_id: str, context: dict) -> dict[str, Any]:
    """
    Save user context to cache and database.

    Args:
        customer_id: The customer's unique identifier.
        context: The context data to save.

    Returns:
        A dictionary indicating success or failure with a message.
    """
    wrapped_fn = tool_wrapper(save_user_context)
    result = await wrapped_fn(customer_id, context)
    if isinstance(result, dict) and result.get("success") is False:
        # Error response from wrapper
        return result

    # result is a tuple (success: bool, message: str)
    success, message = result
    if success:
        return {
            "success": True,
            "message": message
        }
    return {
        "success": False,
        "error_type": "ValidationError",
        "error_message": message,
        "user_message": message
    }


@mcp.tool()
async def retrieve_context(customer_id: str) -> dict[str, Any]:
    """
    Retrieve user context from cache or database.

    Args:
        customer_id: The customer's unique identifier.

    Returns:
        A dictionary containing the user context if found, or an error response.
    """
    wrapped_fn = tool_wrapper(retrieve_user_context)
    result = await wrapped_fn(customer_id)
    if result is None:
        return {
            "success": False,
            "error_type": "NotFound",
            "error_message": f"No context found for customer_id: {customer_id}",
            "user_message": "No context found for this customer."
        }
    if isinstance(result, dict) and result.get("success") is False:
        # Error response from wrapper
        return result
    return {
        "success": True,
        "data": result
    }


# WRAP THE RUN COMMAND IN A FUNCTION
def main():
    """Start the MCP server."""
    logger.info("Context server started.")
    mcp.run()


if __name__ == "__main__":
    main()
