"""Core AiogramMCP server implementation."""

from __future__ import annotations

import asyncio
import logging
from contextlib import suppress

from aiogram import Bot, Dispatcher
from fastmcp import FastMCP

from .context import BotContext
from .tools.broadcast import register_broadcast_tools
from .tools.chats import register_chat_tools
from .tools.messaging import register_messaging_tools
from .tools.users import register_user_tools

logger = logging.getLogger(__name__)


class AiogramMCP:
    """Expose an existing aiogram bot as an MCP server."""

    def __init__(
        self,
        bot: Bot,
        dp: Dispatcher,
        name: str = "aiogram-mcp",
        allowed_chat_ids: list[int] | None = None,
        enable_broadcast: bool = False,
        max_broadcast_recipients: int = 100,
    ) -> None:
        self.bot = bot
        self.dp = dp
        self.name = name
        self.allowed_chat_ids = allowed_chat_ids
        self.enable_broadcast = enable_broadcast
        self.max_broadcast_recipients = max_broadcast_recipients

        self._ctx = BotContext(
            bot=bot,
            dp=dp,
            allowed_chat_ids=allowed_chat_ids,
        )
        self._mcp = FastMCP(
            name=name,
            instructions=(
                "You are connected to a Telegram bot via aiogram-mcp. "
                "You can send messages, inspect users and chats, and perform "
                "administrative actions. Confirm destructive actions first."
            ),
        )

        self._register_tools()
        logger.info("AiogramMCP initialized with server name '%s'", name)

    def _register_tools(self) -> None:
        register_messaging_tools(self._mcp, self._ctx)
        register_user_tools(self._mcp, self._ctx)
        register_chat_tools(self._mcp, self._ctx)

        if self.enable_broadcast:
            register_broadcast_tools(
                self._mcp,
                self._ctx,
                max_recipients=self.max_broadcast_recipients,
            )

    async def run_stdio(self) -> None:
        """Run the MCP server over stdio."""
        logger.info("Starting aiogram-mcp in stdio mode")
        await self._mcp.run_async(transport="stdio")

    async def run_sse(self, host: str = "0.0.0.0", port: int = 8080) -> None:
        """Run the MCP server over SSE."""
        logger.info("Starting aiogram-mcp in SSE mode on %s:%s", host, port)
        await self._mcp.run_async(transport="sse", host=host, port=port)

    async def run_alongside_bot(
        self,
        *,
        transport: str = "stdio",
        host: str = "0.0.0.0",
        port: int = 8080,
    ) -> None:
        """Run bot polling and MCP server concurrently."""
        if transport not in {"stdio", "sse"}:
            raise ValueError("transport must be 'stdio' or 'sse'")

        if transport == "stdio":
            mcp_task = asyncio.create_task(self.run_stdio())
        else:
            mcp_task = asyncio.create_task(self.run_sse(host=host, port=port))

        bot_task = asyncio.create_task(self.dp.start_polling(self.bot))

        try:
            logger.info("Running aiogram polling and MCP server concurrently")
            await asyncio.gather(bot_task, mcp_task)
        finally:
            for task in (bot_task, mcp_task):
                if not task.done():
                    task.cancel()
            with suppress(asyncio.CancelledError):
                await asyncio.gather(bot_task, mcp_task)
            await self.bot.session.close()

    @property
    def fastmcp(self) -> FastMCP:
        """Expose the underlying FastMCP instance for custom tools."""
        return self._mcp
