from contextlib import asynccontextmanager
from functools import partial
from os import getenv
from typing import Any, Dict, List, Optional, Union
from uuid import uuid4

from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from rich import box
from rich.panel import Panel
from starlette.middleware.cors import CORSMiddleware
from starlette.requests import Request

from mielto.agent.agent import Agent
from mielto.os.config import (
    AgentOSConfig,
    DatabaseConfig,
    EvalsConfig,
    EvalsDomainConfig,
    KnowledgeConfig,
    KnowledgeDomainConfig,
    MemoryConfig,
    MemoryDomainConfig,
    MetricsConfig,
    MetricsDomainConfig,
    SessionConfig,
    SessionDomainConfig,
)
from mielto.os.interfaces.base import BaseInterface
from mielto.os.router import get_base_router
from mielto.os.routers.evals import get_eval_router
from mielto.os.routers.knowledge import get_knowledge_router
from mielto.os.routers.memory import get_memory_router
from mielto.os.routers.metrics import get_metrics_router
from mielto.os.routers.session import get_session_router
from mielto.os.settings import MieltoAPISettings
from mielto.os.utils import generate_id
from mielto.team.team import Team
from mielto.workflow.workflow import Workflow


@asynccontextmanager
async def mcp_lifespan(app, mcp_tools):
    """Manage MCP connection lifecycle inside a FastAPI app"""
    # Startup logic: connect to all contextual MCP servers
    for tool in mcp_tools:
        await tool.connect()

    yield

    # Shutdown logic: Close all contextual MCP connections
    for tool in mcp_tools:
        await tool.close()


class AgentOS:
    def __init__(
        self,
        os_id: Optional[str] = None,
        name: Optional[str] = None,
        description: Optional[str] = None,
        version: Optional[str] = None,
        agents: Optional[List[Agent]] = None,
        teams: Optional[List[Team]] = None,
        workflows: Optional[List[Workflow]] = None,
        interfaces: Optional[List[BaseInterface]] = None,
        config: Optional[Union[str, AgentOSConfig]] = None,
        settings: Optional[MieltoAPISettings] = None,
        fastapi_app: Optional[FastAPI] = None,
        lifespan: Optional[Any] = None,
        enable_mcp: bool = False,
        telemetry: bool = True,
    ):
        if not agents and not workflows and not teams:
            raise ValueError("Either agents, teams or workflows must be provided.")

        self.config = self._load_yaml_config(config) if isinstance(config, str) else config

        self.agents: Optional[List[Agent]] = agents
        self.workflows: Optional[List[Workflow]] = workflows
        self.teams: Optional[List[Team]] = teams
        self.interfaces = interfaces or []

        self.settings: MieltoAPISettings = settings or MieltoAPISettings()

        self._app_set = False
        self.fastapi_app: Optional[FastAPI] = None
        if fastapi_app:
            self.fastapi_app = fastapi_app
            self._app_set = True

        self.interfaces = interfaces or []

        self.os_id: Optional[str] = os_id
        self.name = name
        self.version = version
        self.description = description

        self.telemetry = telemetry

        self.enable_mcp = enable_mcp
        self.lifespan = lifespan

        # List of all MCP tools used inside the AgentOS
        self.mcp_tools = []

        if self.agents:
            for agent in self.agents:
                # Track all MCP tools to later handle their connection
                if agent.tools:
                    for tool in agent.tools:
                        # Checking if the tool is a MCPTools or MultiMCPTools instance
                        type_name = type(tool).__name__
                        if type_name in ("MCPTools", "MultiMCPTools"):
                            self.mcp_tools.append(tool)

                agent.initialize_agent()

                # Required for the built-in routes to work
                agent.store_events = True

        if self.teams:
            for team in self.teams:
                # Track all MCP tools to later handle their connection
                if team.tools:
                    for tool in team.tools:
                        # Checking if the tool is a MCPTools or MultiMCPTools instance
                        type_name = type(tool).__name__
                        if type_name in ("MCPTools", "MultiMCPTools"):
                            self.mcp_tools.append(tool)

                team.initialize_team()

                # Required for the built-in routes to work
                team.store_events = True

                for member in team.members:
                    if isinstance(member, Agent):
                        member.team_id = None
                        member.initialize_agent()
                    elif isinstance(member, Team):
                        member.initialize_team()

        if self.workflows:
            for workflow in self.workflows:
                # TODO: track MCP tools in workflow members
                if not workflow.id:
                    workflow.id = generate_id(workflow.name)

        if self.telemetry:
            from mielto.api.os import OSLaunch, log_os_telemetry

            log_os_telemetry(launch=OSLaunch(os_id=self.os_id, data=self._get_telemetry_data()))

    def _make_app(self, lifespan: Optional[Any] = None) -> FastAPI:
        # Adjust the FastAPI app lifespan to handle MCP connections if relevant
        app_lifespan = lifespan
        if self.mcp_tools is not None:
            mcp_tools_lifespan = partial(mcp_lifespan, mcp_tools=self.mcp_tools)
            # If there is already a lifespan, combine it with the MCP lifespan
            if lifespan is not None:
                # Combine both lifespans
                @asynccontextmanager
                async def combined_lifespan(app: FastAPI):
                    # Run both lifespans
                    async with lifespan(app):  # type: ignore
                        async with mcp_tools_lifespan(app):  # type: ignore
                            yield

                app_lifespan = combined_lifespan  # type: ignore
            else:
                app_lifespan = mcp_tools_lifespan

        return FastAPI(
            title=self.name or "Mielto AgentOS",
            version=self.version or "1.0.0",
            description=self.description or "An agent operating system.",
            docs_url="/docs" if self.settings.docs_enabled else None,
            redoc_url="/redoc" if self.settings.docs_enabled else None,
            openapi_url="/openapi.json" if self.settings.docs_enabled else None,
            lifespan=app_lifespan,
        )

    def get_app(self) -> FastAPI:
        if not self.fastapi_app:
            if self.enable_mcp:
                from contextlib import asynccontextmanager

                from mielto.os.mcp import get_mcp_server

                self.mcp_app = get_mcp_server(self)

                final_lifespan = self.mcp_app.lifespan
                if self.lifespan is not None:
                    # Combine both lifespans
                    @asynccontextmanager
                    async def combined_lifespan(app: FastAPI):
                        # Run both lifespans
                        async with self.lifespan(app):  # type: ignore
                            async with self.mcp_app.lifespan(app):  # type: ignore
                                yield

                    final_lifespan = combined_lifespan  # type: ignore

                self.fastapi_app = self._make_app(lifespan=final_lifespan)
            else:
                self.fastapi_app = self._make_app(lifespan=self.lifespan)

        # Add routes
        self.fastapi_app.include_router(get_base_router(self, settings=self.settings))

        for interface in self.interfaces:
            interface_router = interface.get_router()
            self.fastapi_app.include_router(interface_router)

        self._auto_discover_databases()
        self._auto_discover_knowledge_instances()
        self._setup_routers()

        # Mount MCP if needed
        if self.enable_mcp and self.mcp_app:
            self.fastapi_app.mount("/", self.mcp_app)

        # Add middleware (only if app is not set)
        if not self._app_set:

            @self.fastapi_app.exception_handler(HTTPException)
            async def http_exception_handler(request: Request, exc: HTTPException) -> JSONResponse:
                return JSONResponse(
                    status_code=exc.status_code,
                    content={"detail": str(exc.detail)},
                )

            async def general_exception_handler(request: Request, call_next):
                try:
                    return await call_next(request)
                except Exception as e:
                    return JSONResponse(
                        status_code=e.status_code if hasattr(e, "status_code") else 500,  # type: ignore
                        content={"detail": str(e)},
                    )

            self.fastapi_app.middleware("http")(general_exception_handler)

            self.fastapi_app.add_middleware(
                CORSMiddleware,
                allow_origins=self.settings.cors_origin_list,  # type: ignore
                allow_credentials=True,
                allow_methods=["*"],
                allow_headers=["*"],
                expose_headers=["*"],
            )

        return self.fastapi_app

    def get_routes(self) -> List[Any]:
        """Retrieve all routes from the FastAPI app.

        Returns:
            List[Any]: List of routes included in the FastAPI app.
        """
        app = self.get_app()

        return app.routes

    def _get_telemetry_data(self) -> Dict[str, Any]:
        """Get the telemetry data for the OS"""
        return {
            "agents": [agent.id for agent in self.agents] if self.agents else None,
            "teams": [team.id for team in self.teams] if self.teams else None,
            "workflows": [workflow.id for workflow in self.workflows] if self.workflows else None,
            "interfaces": [interface.type for interface in self.interfaces] if self.interfaces else None,
        }

    def _load_yaml_config(self, config_file_path: str) -> AgentOSConfig:
        """Load a YAML config file and return the configuration as an AgentOSConfig instance."""
        from pathlib import Path

        import yaml

        # Validate that the path points to a YAML file
        path = Path(config_file_path)
        if path.suffix.lower() not in [".yaml", ".yml"]:
            raise ValueError(f"Config file must have a .yaml or .yml extension, got: {config_file_path}")

        # Load the YAML file
        with open(config_file_path, "r") as f:
            return AgentOSConfig.model_validate(yaml.safe_load(f))

    def _auto_discover_databases(self) -> None:
        """Auto-discover the databases used by all contextual agents, teams and workflows."""
        dbs = {}

        for agent in self.agents or []:
            if agent.db:
                dbs[agent.db.id] = agent.db
            if agent.knowledge and agent.knowledge.contents_db:
                dbs[agent.knowledge.contents_db.id] = agent.knowledge.contents_db

        for team in self.teams or []:
            if team.db:
                dbs[team.db.id] = team.db
            if team.knowledge and team.knowledge.contents_db:
                dbs[team.knowledge.contents_db.id] = team.knowledge.contents_db

        for workflow in self.workflows or []:
            if workflow.db:
                dbs[workflow.db.id] = workflow.db

        for interface in self.interfaces or []:
            if interface.agent and interface.agent.db:
                dbs[interface.agent.db.id] = interface.agent.db
            elif interface.team and interface.team.db:
                dbs[interface.team.db.id] = interface.team.db

        self.dbs = dbs

    def _auto_discover_knowledge_instances(self) -> None:
        """Auto-discover the knowledge instances used by all contextual agents, teams and workflows."""
        knowledge_instances = []
        for agent in self.agents or []:
            if agent.knowledge:
                knowledge_instances.append(agent.knowledge)

        for team in self.teams or []:
            if team.knowledge:
                knowledge_instances.append(team.knowledge)

        self.knowledge_instances = knowledge_instances

    def _get_session_config(self) -> SessionConfig:
        session_config = self.config.session if self.config and self.config.session else SessionConfig()

        if session_config.dbs is None:
            session_config.dbs = []

        multiple_dbs: bool = len(self.dbs.keys()) > 1
        dbs_with_specific_config = [db.db_id for db in session_config.dbs]

        for db_id in self.dbs.keys():
            if db_id not in dbs_with_specific_config:
                session_config.dbs.append(
                    DatabaseConfig(
                        db_id=db_id,
                        domain_config=SessionDomainConfig(
                            display_name="Sessions" if not multiple_dbs else "Sessions in database '" + db_id + "'"
                        ),
                    )
                )

        return session_config

    def _get_memory_config(self) -> MemoryConfig:
        memory_config = self.config.memory if self.config and self.config.memory else MemoryConfig()

        if memory_config.dbs is None:
            memory_config.dbs = []

        multiple_dbs: bool = len(self.dbs.keys()) > 1
        dbs_with_specific_config = [db.db_id for db in memory_config.dbs]

        for db_id in self.dbs.keys():
            if db_id not in dbs_with_specific_config:
                memory_config.dbs.append(
                    DatabaseConfig(
                        db_id=db_id,
                        domain_config=MemoryDomainConfig(
                            display_name="Memory" if not multiple_dbs else "Memory in database '" + db_id + "'"
                        ),
                    )
                )

        return memory_config

    def _get_knowledge_config(self) -> KnowledgeConfig:
        knowledge_config = self.config.knowledge if self.config and self.config.knowledge else KnowledgeConfig()

        if knowledge_config.dbs is None:
            knowledge_config.dbs = []

        multiple_dbs: bool = len(self.dbs.keys()) > 1
        dbs_with_specific_config = [db.db_id for db in knowledge_config.dbs]

        for db_id in self.dbs.keys():
            if db_id not in dbs_with_specific_config:
                knowledge_config.dbs.append(
                    DatabaseConfig(
                        db_id=db_id,
                        domain_config=KnowledgeDomainConfig(
                            display_name="Knowledge" if not multiple_dbs else "Knowledge in database " + db_id
                        ),
                    )
                )

        return knowledge_config

    def _get_metrics_config(self) -> MetricsConfig:
        metrics_config = self.config.metrics if self.config and self.config.metrics else MetricsConfig()

        if metrics_config.dbs is None:
            metrics_config.dbs = []

        multiple_dbs: bool = len(self.dbs.keys()) > 1
        dbs_with_specific_config = [db.db_id for db in metrics_config.dbs]

        for db_id in self.dbs.keys():
            if db_id not in dbs_with_specific_config:
                metrics_config.dbs.append(
                    DatabaseConfig(
                        db_id=db_id,
                        domain_config=MetricsDomainConfig(
                            display_name="Metrics" if not multiple_dbs else "Metrics in database '" + db_id + "'"
                        ),
                    )
                )

        return metrics_config

    def _get_evals_config(self) -> EvalsConfig:
        evals_config = self.config.evals if self.config and self.config.evals else EvalsConfig()

        if evals_config.dbs is None:
            evals_config.dbs = []

        multiple_dbs: bool = len(self.dbs.keys()) > 1
        dbs_with_specific_config = [db.db_id for db in evals_config.dbs]

        for db_id in self.dbs.keys():
            if db_id not in dbs_with_specific_config:
                evals_config.dbs.append(
                    DatabaseConfig(
                        db_id=db_id,
                        domain_config=EvalsDomainConfig(
                            display_name="Evals" if not multiple_dbs else "Evals in database '" + db_id + "'"
                        ),
                    )
                )

        return evals_config

    def _setup_routers(self) -> None:
        """Add all routers to the FastAPI app."""
        if not self.dbs or not self.fastapi_app:
            return

        routers = [
            get_session_router(dbs=self.dbs),
            get_memory_router(dbs=self.dbs),
            get_eval_router(dbs=self.dbs, agents=self.agents, teams=self.teams),
            get_metrics_router(dbs=self.dbs),
            get_knowledge_router(knowledge_instances=self.knowledge_instances),
        ]

        for router in routers:
            self.fastapi_app.include_router(router)

    def set_os_id(self) -> str:
        # If os_id is already set, keep it instead of overriding with UUID
        if self.os_id is None:
            self.os_id = str(uuid4())

        return self.os_id

    def serve(
        self,
        app: Union[str, FastAPI],
        *,
        host: str = "localhost",
        port: int = 7777,
        reload: bool = False,
        workers: Optional[int] = None,
        **kwargs,
    ):
        import uvicorn

        if getenv("MIELTO_API_RUNTIME", "").lower() == "stg":
            public_endpoint = "https://os-stg.mielto.com/"
        else:
            public_endpoint = "https://os.mielto.com/"

        # Create a terminal panel to announce OS initialization and provide useful info
        from rich.align import Align
        from rich.console import Console, Group

        aligned_endpoint = Align.center(f"[bold cyan]{public_endpoint}[/bold cyan]")
        connection_endpoint = f"\n\n[bold dark_orange]Running on:[/bold dark_orange] http://{host}:{port}"

        console = Console()
        console.print(
            Panel(
                Group(aligned_endpoint, connection_endpoint),
                title="AgentOS",
                expand=False,
                border_style="dark_orange",
                box=box.DOUBLE_EDGE,
                padding=(2, 2),
            )
        )

        uvicorn.run(app=app, host=host, port=port, reload=reload, workers=workers, **kwargs)
