from __future__ import annotations

import logging
from importlib import resources
from typing import Any

from amsdal_server.apps.common.events.server import MiddlewareSetupContext
from amsdal_server.apps.common.events.server import RouterSetupContext
from amsdal_utils.events import AsyncNextFn
from amsdal_utils.events import EventListener
from amsdal_utils.events import NextFn
from starlette.staticfiles import StaticFiles

from amsdal_ml.ml_config import ml_config

logger = logging.getLogger(__name__)


class MCPRouterSetupListener(EventListener[RouterSetupContext]):
    def handle(
        self,
        context: RouterSetupContext,
        next_fn: NextFn[RouterSetupContext],
    ) -> RouterSetupContext:
        next_context = self._prepare_context(context)
        return next_fn(next_context)

    async def ahandle(
        self,
        context: RouterSetupContext,
        next_fn: AsyncNextFn[RouterSetupContext],
    ) -> RouterSetupContext:
        next_context = self._prepare_context(context)
        return await next_fn(next_context)

    def _prepare_context(self, context: RouterSetupContext) -> RouterSetupContext:
        from amsdal_ml.mcp_server.server_model_explorer import mcp_server

        logger.info('Mounting MCP to FastAPI app...')
        mcp_app: Any = mcp_server.sse_app()
        context.app.mount('/mcp', mcp_app)
        logger.info('MCP mounted at /mcp')

        if ml_config.oauth_enabled:
            from amsdal_ml.mcp_server.mcp_oauth import init_oauth_routes
            from amsdal_ml.mcp_server.mcp_oauth import oauth_router

            init_oauth_routes()
            static_dir = resources.files('amsdal_ml.mcp_server.mcp_oauth').joinpath('static')
            context.app.mount('/auth/static', StaticFiles(directory=str(static_dir)), name='oauth-static')
            context.app.include_router(oauth_router)
            logger.info('OAuth routes mounted')

        return context


class MCPMiddlewareSetupListener(EventListener[MiddlewareSetupContext]):
    def handle(
        self,
        context: MiddlewareSetupContext,
        next_fn: NextFn[MiddlewareSetupContext],
    ) -> MiddlewareSetupContext:
        next_context = self._prepare_context(context)
        return next_fn(next_context)

    async def ahandle(
        self,
        context: MiddlewareSetupContext,
        next_fn: AsyncNextFn[MiddlewareSetupContext],
    ) -> MiddlewareSetupContext:
        next_context = self._prepare_context(context)
        return await next_fn(next_context)

    def _prepare_context(self, context: MiddlewareSetupContext) -> MiddlewareSetupContext:
        if ml_config.oauth_enabled:
            from amsdal_ml.mcp_server.mcp_oauth.middlwares.www_authenticate import install_www_authenticate_middleware

            install_www_authenticate_middleware(context.app)
            logger.info('OAuth WWW-Authenticate middleware installed')

        return context
