import asyncio
import contextlib
import os
import typing

import uvicorn
from fastmcp import FastMCP
from starlette.applications import Starlette
from starlette.routing import Mount

from src.services.telemetry_service import TelemetryService
from src.tools import all_tools
from src.tools.annotations import get_tool_annotations, get_tool_tags, validate_tool_annotation_coverage
from src.utils.config import settings
from src.utils.error import agent_hint_handler
from src.utils.logger import configure_logging, get_logger
from src.utils.telemetry import set_telemetry_service, wrap_tool_with_telemetry
from src.version import __version__

# Configure logging early
configure_logging(
    log_level=settings.LOG_LEVEL, log_format=settings.LOG_FORMAT, force_stderr=(settings.MCP_MODE == "stdio")
)
logger = get_logger("lucius-mcp")
telemetry_service = TelemetryService()
set_telemetry_service(telemetry_service)


# Initialize FastMCP server
mcp = FastMCP(
    name="lucius-mcp",
    version=__version__,
)

# Register tools
validate_tool_annotation_coverage({tool.__name__ for tool in all_tools})
for tool in all_tools:
    mcp.tool(tags=get_tool_tags(tool.__name__), annotations=get_tool_annotations(tool.__name__))(
        wrap_tool_with_telemetry(tool)
    )

# The ASGI app and main app are created lazily or only when needed for HTTP mode
_mcp_asgi = None


def get_mcp_asgi() -> Starlette:
    global _mcp_asgi
    if _mcp_asgi is None:
        _mcp_asgi = mcp.http_app()
    return _mcp_asgi


@contextlib.asynccontextmanager
async def lifespan(app: Starlette) -> typing.AsyncGenerator[None]:
    """
    Lifespan context manager for Starlette application.
    Handles startup and shutdown events.
    """
    telemetry_service.log_status()
    telemetry_service.emit_startup_event()
    logger.info(f"Starting Lucius MCP Server in {settings.MCP_MODE} mode")
    mcp_asgi = get_mcp_asgi()
    # Ensure MCP task group is initialized by entering its lifespan
    if hasattr(mcp_asgi, "lifespan"):
        async with mcp_asgi.lifespan(app):
            yield
    else:
        yield
    logger.info("Shutting down Lucius MCP Server")


# Create main Starlette application lazily
def get_app() -> Starlette | None:
    if settings.MCP_MODE == "http":
        return Starlette(
            debug=False,
            lifespan=lifespan,
            exception_handlers={Exception: agent_hint_handler},
            routes=[
                # Mount the FastMCP ASGI app under /
                Mount("/", app=get_mcp_asgi())
            ],
        )
    else:
        return None


# For uvicorn.run("src.main:app", ...)
app = get_app()


async def _run_stdio() -> None:
    telemetry_service.log_status()
    telemetry_service.emit_startup_event()
    await mcp.run_stdio_async(show_banner=False, log_level=settings.LOG_LEVEL)


def start() -> None:
    """Entry point for running the application directly."""
    if settings.MCP_MODE == "stdio":
        try:
            asyncio.run(_run_stdio())
        except KeyboardInterrupt:
            os._exit(0)
    elif settings.MCP_MODE == "http":
        uvicorn.run("src.main:app", host=settings.HOST, port=settings.PORT, reload=True, ws="wsproto")
    else:
        raise ValueError(f"Invalid MCP mode: {settings.MCP_MODE}")


if __name__ == "__main__":
    start()
