"""Tool registration for Meta Ads MCP Server.

Provides function to register all tools with FastMCP server.
"""
from typing import TYPE_CHECKING, Dict, Tuple, Any
import asyncio
import json

from mcp import types as mcp_types
from mcp.server.fastmcp import FastMCP

from .base import ToolHandler

# Import all tool submodules
from . import (
    ad_management,
    advantage,
    aggregation,
    analysis,
    analysis_advanced,
    ar_vr,
    attribution,
    audiences,
    base,
    bidding,
    brand_safety,
    business,
    catalog,
    commerce,
    comparison,
    competitive,
    conversion,
    cost_efficiency,
    creative,
    diagnostics,
    finance,
    fraud,
    insights,
    leads,
    monitoring,
    multi_account,
    multi_platform,
    optimization,
    pixels,
    portfolio,
    predictive,
    realtime,
    reels,
    reporting,
    roas,
    scale,
    scenario,
    scheduling,
    search,
    strategy,
    tactical,
    targeting,
    validation,
    video,
    video_mastery,
    messaging,
    partnerships,
    experiments,
    threads,
    ccco,
    signals,
    whatsapp_flows,
)

if TYPE_CHECKING:
    from ..meta_client import MetaClient

# All tool modules
TOOL_MODULES = [
    ad_management,
    advantage,
    aggregation,
    analysis,
    analysis_advanced,
    ar_vr,
    attribution,
    audiences,
    base,
    bidding,
    brand_safety,
    business,
    catalog,
    commerce,
    comparison,
    competitive,
    conversion,
    cost_efficiency,
    creative,
    diagnostics,
    finance,
    fraud,
    insights,
    leads,
    monitoring,
    multi_account,
    multi_platform,
    optimization,
    pixels,
    portfolio,
    predictive,
    realtime,
    reels,
    reporting,
    roas,
    scale,
    scenario,
    scheduling,
    search,
    strategy,
    tactical,
    targeting,
    validation,
    video,
    video_mastery,
    messaging,
    partnerships,
    experiments,
    threads,
    ccco,
    signals,
    whatsapp_flows,
]


def _json_safe(payload: Any) -> Any:
    """Convert tool payloads to JSON-serializable objects."""
    return json.loads(json.dumps(payload, default=str))


def register_all_tools(mcp: FastMCP) -> None:
    """Register all tools with the FastMCP server.
    
    This function collects tools from all modules and registers them
    using FastMCP's low-level API for backward compatibility with
    existing tool definitions.
    """
    from ..meta_client import get_client
    from ..logging import get_logger
    from ..auth.request_context import bind_user_id
    from ..error_handlers import get_error_handler
    from .. import licensing as _licensing
    
    logger = get_logger(__name__)
    error_handler = get_error_handler()
    
    # Get client (will be available after lifespan starts)
    def get_client_lazy():
        return get_client()
    
    all_tools: Dict[str, Tuple[mcp_types.Tool, ToolHandler]] = {}
    tools_loaded_with_client = False
    tools_load_lock = asyncio.Lock()

    def _reset_tools_loaded():
        """Reset the tools loaded flag (for testing only)."""
        nonlocal tools_loaded_with_client
        tools_loaded_with_client = False

    # Expose for test teardown
    mcp._tools_reset = _reset_tools_loaded

    def load_tools(client: Any) -> None:
        nonlocal tools_loaded_with_client
        loaded_tools: Dict[str, Tuple[mcp_types.Tool, ToolHandler]] = {}
        for module in TOOL_MODULES:
            try:
                tools = module.get_tools(client)
                loaded_tools.update(tools)
            except Exception as exc:
                logger.error("tool_load_error module=%s error=%s", module.__name__, exc)
        all_tools.clear()
        all_tools.update(loaded_tools)
        tools_loaded_with_client = True

    async def ensure_tools_loaded_with_client() -> None:
        nonlocal tools_loaded_with_client
        if tools_loaded_with_client:
            return
        async with tools_load_lock:
            if tools_loaded_with_client:
                return
            load_tools(get_client_lazy())

    def resolve_authenticated_user_id() -> str | None:
        """Read authenticated user id from FastMCP auth context when present."""
        try:
            from mcp.server.auth.middleware.auth_context import get_access_token

            token = get_access_token()
            if token and token.client_id:
                return token.client_id
        except Exception:
            # No auth middleware context in stdio or unauthenticated environments.
            return None
        return None
    
    # Collect tools from all modules
    # Note: We defer client access until tool execution
    for module in TOOL_MODULES:
        try:
            # Pass None for now - handlers will use get_client() internally
            tools = module.get_tools(None)  # type: ignore
            all_tools.update(tools)
        except Exception as exc:
            logger.warning("deferred_loading module=%s error=%s", module.__name__, exc)
    
    # Register each tool with FastMCP using low-level API
    @mcp._mcp_server.list_tools()
    async def handle_list_tools() -> list[mcp_types.Tool]:
        """List all available tools, filtered by license tier."""
        # Reload tools with actual client
        try:
            await ensure_tools_loaded_with_client()
        except RuntimeError:
            # Client not ready yet, return deferred registration.
            pass
        tier = _licensing.get_license_tier()
        return [tool for tool, _ in all_tools.values()
                if _licensing.is_tool_allowed(tool.name, tier)]

    @mcp._mcp_server.call_tool()
    async def handle_call_tool(
        name: str, arguments: dict | None
    ) -> mcp_types.CallToolResult | dict[str, Any]:
        """Handle tool execution with structured output."""
        try:
            await ensure_tools_loaded_with_client()

            if name not in all_tools:
                raise ValueError(f"Unknown tool: {name}")

            tier = _licensing.get_license_tier()
            if not _licensing.is_tool_allowed(name, tier):
                return mcp_types.CallToolResult(
                    content=[mcp_types.TextContent(
                        type="text",
                        text=json.dumps({
                            "error": {
                                "code": "LICENSE_REQUIRED",
                                "message": (
                                    f"Tool '{name}' requires a paid license. "
                                    f"Current tier: {tier}. "
                                    "Contact the developer for a license key."
                                ),
                            }
                        }, indent=2),
                    )],
                    isError=True,
                )

            _, handler = all_tools[name]
            user_id = resolve_authenticated_user_id()
            with bind_user_id(user_id):
                result = await handler(arguments or {})
            
            if not result.ok:
                return mcp_types.CallToolResult(
                    content=[
                        mcp_types.TextContent(
                            type="text",
                            text=json.dumps({"error": result.error}, indent=2, default=str),
                        )
                    ],
                    structuredContent={"error": result.error},
                    isError=True,
                )

            if isinstance(result.data, dict):
                return _json_safe(result.data)

            return {"result": _json_safe(result.data)}
        except Exception as e:
            logger.error("tool_execution_error tool=%s error=%s", name, e)
            handled_error = await error_handler.handle(e, context={"tool": name})
            if handled_error.ok:
                if isinstance(handled_error.data, dict):
                    return _json_safe(handled_error.data)
                return {"result": _json_safe(handled_error.data)}
            error_payload = _json_safe(
                handled_error.error
                if handled_error.error is not None
                else {"code": "INTERNAL_ERROR", "message": str(e)}
            )
            return mcp_types.CallToolResult(
                content=[
                    mcp_types.TextContent(
                        type="text",
                        text=json.dumps({"error": error_payload}, indent=2, default=str),
                    )
                ],
                structuredContent={"error": error_payload},
                isError=True,
            )


# Legacy function for backward compatibility
def register_tools(server: Any, client: "MetaClient") -> None:
    """Legacy registration function - deprecated.
    
    Use register_all_tools(mcp) with FastMCP instead.
    """
    import warnings
    warnings.warn(
        "register_tools() is deprecated. Use register_all_tools() with FastMCP.",
        DeprecationWarning,
        stacklevel=2
    )
    
    all_tools: Dict[str, Tuple[mcp_types.Tool, ToolHandler]] = {}
    
    for module in TOOL_MODULES:
        try:
            tools = module.get_tools(client)
            all_tools.update(tools)
        except Exception as e:
            print(f"Error loading tools from module {module.__name__}: {e}")

    @server.list_tools()
    async def handle_list_tools() -> list[mcp_types.Tool]:
        return [tool for tool, _ in all_tools.values()]

    @server.call_tool()
    async def handle_call_tool(
        name: str, arguments: dict | None
    ) -> mcp_types.CallToolResult | dict[str, Any]:
        if name not in all_tools:
            raise ValueError(f"Unknown tool: {name}")

        _, handler = all_tools[name]
        
        try:
            result = await handler(arguments or {})
            
            if not result.ok:
                return mcp_types.CallToolResult(
                    content=[
                        mcp_types.TextContent(
                            type="text",
                            text=json.dumps({"error": result.error}, indent=2, default=str),
                        )
                    ],
                    structuredContent={"error": result.error},
                    isError=True,
                )

            if isinstance(result.data, dict):
                return _json_safe(result.data)
            return {"result": _json_safe(result.data)}
        except Exception as e:
            return mcp_types.CallToolResult(
                content=[mcp_types.TextContent(type="text", text=f"Tool execution error: {str(e)}")],
                structuredContent={"error": {"message": str(e)}},
                isError=True,
            )
