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

from fastapi import APIRouter, FastAPI, HTTPException
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse
from fastapi.routing import APIRoute
from httpx import HTTPStatusError
from rich import box
from rich.panel import Panel
from starlette.requests import Request

from agno.agent import Agent, RemoteAgent
from agno.db.base import AsyncBaseDb, BaseDb
from agno.knowledge.knowledge import Knowledge
from agno.os.config import (
    AgentOSConfig,
    AuthorizationConfig,
    DatabaseConfig,
    EvalsConfig,
    EvalsDomainConfig,
    KnowledgeConfig,
    KnowledgeDatabaseConfig,
    KnowledgeDomainConfig,
    KnowledgeInstanceConfig,
    MemoryConfig,
    MemoryDomainConfig,
    MetricsConfig,
    MetricsDomainConfig,
    SessionConfig,
    SessionDomainConfig,
    TracesConfig,
    TracesDomainConfig,
)
from agno.os.interfaces.base import BaseInterface
from agno.os.router import get_base_router, get_websocket_router
from agno.os.routers.agents import get_agent_router
from agno.os.routers.approvals import get_approval_router
from agno.os.routers.components import get_components_router
from agno.os.routers.database import get_database_router
from agno.os.routers.evals import get_eval_router
from agno.os.routers.health import get_health_router
from agno.os.routers.home import get_home_router
from agno.os.routers.knowledge import get_knowledge_router
from agno.os.routers.memory import get_memory_router
from agno.os.routers.metrics import get_metrics_router
from agno.os.routers.registry import get_registry_router
from agno.os.routers.schedules import get_schedule_router
from agno.os.routers.session import get_session_router
from agno.os.routers.teams import get_team_router
from agno.os.routers.traces import get_traces_router
from agno.os.routers.workflows import get_workflow_router
from agno.os.settings import AgnoAPISettings
from agno.os.utils import (
    _generate_knowledge_id,
    collect_mcp_tools_from_team,
    collect_mcp_tools_from_workflow,
    find_conflicting_routes,
    load_yaml_config,
    resolve_origins,
    setup_tracing_for_os,
    update_cors_middleware,
)
from agno.registry import Registry
from agno.remote.base import RemoteDb, RemoteKnowledge
from agno.team import RemoteTeam, Team
from agno.utils.log import log_debug, log_error, log_info, log_warning
from agno.utils.string import generate_id, generate_id_from_name
from agno.workflow import RemoteWorkflow, Workflow


@asynccontextmanager
async def mcp_lifespan(_, mcp_tools):
    """Manage MCP connection lifecycle inside a FastAPI app"""
    for tool in mcp_tools:
        await tool.connect()

    yield

    for tool in mcp_tools:
        await tool.close()


@asynccontextmanager
async def http_client_lifespan(_):
    """Manage httpx client lifecycle for proper connection pool cleanup."""
    from agno.utils.http import aclose_default_clients

    yield

    await aclose_default_clients()


@asynccontextmanager
async def db_lifespan(app: FastAPI, agent_os: "AgentOS"):
    """Initializes databases in the event loop and closes them on shutdown."""
    if agent_os.auto_provision_dbs:
        agent_os._initialize_sync_databases()
        await agent_os._initialize_async_databases()

    yield

    await agent_os._close_databases()


@asynccontextmanager
async def scheduler_lifespan(app: FastAPI, agent_os: "AgentOS"):
    """Start and stop the scheduler poller."""
    from agno.scheduler import ScheduleExecutor, SchedulePoller

    if agent_os._scheduler_base_url is None:
        log_info(
            "scheduler_base_url not set, using default http://127.0.0.1:7777. "
            "If your server is running on a different port, set scheduler_base_url to match."
        )
    base_url = agent_os._scheduler_base_url or "http://127.0.0.1:7777"
    internal_token = agent_os._internal_service_token
    if internal_token is None:
        raise ValueError("internal_service_token must be set when scheduler is enabled")

    executor = ScheduleExecutor(
        base_url=base_url,
        internal_service_token=internal_token,
    )
    poller = SchedulePoller(
        db=agent_os.db,
        executor=executor,
        poll_interval=agent_os._scheduler_poll_interval,
    )

    app.state.scheduler_executor = executor
    app.state.scheduler_poller = poller
    await poller.start()

    yield

    await poller.stop()


def _combine_app_lifespans(lifespans: list) -> Any:
    """Combine multiple FastAPI app lifespan context managers into one."""
    if len(lifespans) == 1:
        return lifespans[0]

    from contextlib import asynccontextmanager

    @asynccontextmanager
    async def combined_lifespan(app):
        async def _run_nested(index: int):
            if index >= len(lifespans):
                yield
                return

            async with lifespans[index](app):
                async for _ in _run_nested(index + 1):
                    yield

        async for _ in _run_nested(0):
            yield

    return combined_lifespan


def _get_disabled_feature_router(prefix: str, tag: str, requires: str) -> APIRouter:
    """Return a stub router that returns 503 for a feature that requires a missing dependency."""
    detail = f"{tag} not available: pass a `{requires}` to AgentOS to enable this feature."
    router = APIRouter(tags=[tag])
    for path in [prefix, f"{prefix}/{{path:path}}"]:

        @router.api_route(path, methods=["GET", "POST", "PUT", "PATCH", "DELETE"])
        async def _disabled() -> None:
            raise HTTPException(status_code=503, detail=detail)

    return router


class AgentOS:
    def __init__(
        self,
        id: Optional[str] = None,
        name: Optional[str] = None,
        description: Optional[str] = None,
        version: Optional[str] = None,
        db: Optional[Union[BaseDb, AsyncBaseDb]] = None,
        agents: Optional[List[Union[Agent, RemoteAgent]]] = None,
        teams: Optional[List[Union[Team, RemoteTeam]]] = None,
        workflows: Optional[List[Union[Workflow, RemoteWorkflow]]] = None,
        knowledge: Optional[List[Knowledge]] = None,
        interfaces: Optional[List[BaseInterface]] = None,
        a2a_interface: bool = False,
        authorization: bool = False,
        authorization_config: Optional[AuthorizationConfig] = None,
        cors_allowed_origins: Optional[List[str]] = None,
        config: Optional[Union[str, AgentOSConfig]] = None,
        settings: Optional[AgnoAPISettings] = None,
        lifespan: Optional[Any] = None,
        enable_mcp_server: bool = False,
        base_app: Optional[FastAPI] = None,
        on_route_conflict: Literal["preserve_agentos", "preserve_base_app", "error"] = "preserve_agentos",
        tracing: bool = False,
        auto_provision_dbs: bool = True,
        run_hooks_in_background: bool = False,
        telemetry: bool = True,
        registry: Optional[Registry] = None,
        scheduler: bool = False,
        scheduler_poll_interval: int = 15,
        scheduler_base_url: Optional[str] = None,
        internal_service_token: Optional[str] = None,
    ):
        """Initialize AgentOS.

        Args:
            id: Unique identifier for this AgentOS instance
            name: Name of the AgentOS instance
            description: Description of the AgentOS instance
            version: Version of the AgentOS instance
            db: Default database for the AgentOS instance. Agents, teams and workflows with no db will use this one.
            agents: List of agents to include in the OS
            teams: List of teams to include in the OS
            workflows: List of workflows to include in the OS
            knowledge: List of knowledge bases to include in the OS
            interfaces: List of interfaces to include in the OS
            a2a_interface: Whether to expose the OS agents and teams in an A2A server
            config: Configuration file path or AgentOSConfig instance
            settings: API settings for the OS
            lifespan: Optional lifespan context manager for the FastAPI app
            enable_mcp_server: Whether to enable MCP (Model Context Protocol)
            base_app: Optional base FastAPI app to use for the AgentOS. All routes and middleware will be added to this app.
            on_route_conflict: What to do when a route conflict is detected in case a custom base_app is provided.
            auto_provision_dbs: Whether to automatically provision databases
            authorization: Whether to enable authorization
            authorization_config: Configuration for the authorization middleware
            cors_allowed_origins: List of allowed CORS origins (will be merged with default Agno domains)
            tracing: If True, enables OpenTelemetry tracing for all agents and teams in the OS
            run_hooks_in_background: If True, run agent/team pre/post hooks as FastAPI background tasks (non-blocking)
            telemetry: Whether to enable telemetry
            registry: Optional registry to use for the AgentOS
            scheduler: Whether to enable the cron scheduler
            scheduler_poll_interval: Seconds between scheduler poll cycles (default: 15)
            scheduler_base_url: Base URL for scheduler HTTP calls (default: http://127.0.0.1:7777)
            internal_service_token: Token for scheduler-to-OS auth (auto-generated if not provided)

        """
        if not agents and not workflows and not teams and not knowledge and not db:
            raise ValueError("Either agents, teams, workflows, knowledge bases or a database must be provided.")

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

        self.agents: Optional[List[Union[Agent, RemoteAgent]]] = agents
        self.workflows: Optional[List[Union[Workflow, RemoteWorkflow]]] = workflows
        self.teams: Optional[List[Union[Team, RemoteTeam]]] = teams
        self.a2a_interface = a2a_interface
        self.knowledge = knowledge
        self.settings: AgnoAPISettings = settings or AgnoAPISettings()
        self.auto_provision_dbs = auto_provision_dbs
        self._app_set = False

        if base_app:
            self.base_app: Optional[FastAPI] = base_app
            self._app_set = True
            self.on_route_conflict = on_route_conflict
        else:
            self.base_app = None
            self._app_set = False
            self.on_route_conflict = on_route_conflict

        self.interfaces = interfaces or []

        self.name = name

        self.id = id
        if not self.id:
            self.id = generate_id(self.name) if self.name else str(uuid4())

        self.version = version
        self.description = description
        self.db = db

        self.telemetry = telemetry
        self.tracing = tracing

        self.enable_mcp_server = enable_mcp_server
        self.lifespan = lifespan

        self.registry = registry

        # RBAC
        self.authorization = authorization
        self.authorization_config = authorization_config

        # CORS configuration - merge user-provided origins with defaults from settings
        self.cors_allowed_origins = resolve_origins(cors_allowed_origins, self.settings.cors_origin_list)

        # If True, run agent/team hooks as FastAPI background tasks
        self.run_hooks_in_background = run_hooks_in_background

        # Scheduler configuration
        self._scheduler_enabled = scheduler
        self._scheduler_poll_interval = scheduler_poll_interval
        self._scheduler_base_url = scheduler_base_url
        if self._scheduler_enabled and not internal_service_token:
            import secrets

            internal_service_token = secrets.token_urlsafe(32)
        self._internal_service_token = internal_service_token

        # List of all MCP tools used inside the AgentOS
        self.mcp_tools: List[Any] = []
        self._mcp_app: Optional[Any] = None

        self._initialize_agents()
        self._initialize_teams()
        self._initialize_workflows()

        # Populate registry with code-defined agents/teams
        self._populate_registry()

        # Check for duplicate IDs
        self._raise_if_duplicate_ids()

        if self.tracing:
            self._setup_tracing()

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

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

    def _add_agent_os_to_lifespan_function(self, lifespan):
        """
        Inspect a lifespan function and wrap it to pass agent_os if it accepts it.

        Returns:
            A wrapped lifespan that passes agent_os if the lifespan function expects it.
        """
        # Getting the actual function inside the lifespan
        lifespan_function = lifespan
        if hasattr(lifespan, "__wrapped__"):
            lifespan_function = lifespan.__wrapped__

        try:
            from inspect import signature

            # Inspecting the lifespan function signature to find its parameters
            sig = signature(lifespan_function)
            params = list(sig.parameters.keys())

            # If the lifespan function expects the 'agent_os' parameter, add it
            if "agent_os" in params:
                return partial(lifespan, agent_os=self)
            else:
                return lifespan

        except (ValueError, TypeError):
            return lifespan

    def resync(self, app: FastAPI) -> None:
        """Resync the AgentOS to discover, initialize and configure: agents, teams, workflows, databases and knowledge bases."""
        self._initialize_agents()
        self._initialize_teams()
        self._initialize_workflows()

        # Populate registry with code-defined agents/teams
        self._populate_registry()

        # Check for duplicate IDs
        self._raise_if_duplicate_ids()
        self._auto_discover_databases()
        self._auto_discover_knowledge_instances()

        if self.enable_mcp_server:
            from agno.os.mcp import get_mcp_server

            self._mcp_app = get_mcp_server(self)

        self._reprovision_routers(app=app)

    def _reprovision_routers(self, app: FastAPI) -> None:
        """Re-provision all routes for the AgentOS."""
        updated_routers = [
            get_home_router(self),
            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),
            get_traces_router(dbs=self.dbs),
            get_database_router(self, settings=self.settings),
        ]
        # Routes that require a database
        if self.db is not None:
            if isinstance(self.db, BaseDb):
                updated_routers.append(get_components_router(os_db=self.db, registry=self.registry))
            else:
                updated_routers.append(_get_disabled_feature_router("/components", "Components", "sync db (BaseDb)"))
            updated_routers.append(get_schedule_router(os_db=self.db, settings=self.settings))
            updated_routers.append(get_approval_router(os_db=self.db, settings=self.settings))
        else:
            for prefix, tag in [("/components", "Components"), ("/schedules", "Schedules"), ("/approvals", "Approvals")]:
                updated_routers.append(_get_disabled_feature_router(prefix, tag, "db"))
        # Registry router
        if self.registry is not None:
            updated_routers.append(get_registry_router(registry=self.registry))
        else:
            updated_routers.append(_get_disabled_feature_router("/registry", "Registry", "registry"))

        # Clear all previously existing routes
        app.router.routes = [
            route
            for route in app.router.routes
            if hasattr(route, "path")
            and route.path in ["/docs", "/redoc", "/openapi.json", "/docs/oauth2-redirect"]
            or route.path.startswith("/mcp")  # type: ignore
        ]

        # Add the built-in routes
        self._add_built_in_routes(app=app)

        # Add the updated routes
        for router in updated_routers:
            self._add_router(app, router)

        # Mount MCP if needed
        if self.enable_mcp_server and self._mcp_app:
            app.mount("/", self._mcp_app)

    def _add_built_in_routes(self, app: FastAPI) -> None:
        """Add all AgentOSbuilt-in routes to the given app."""
        # Add the home router if MCP server is not enabled
        if not self.enable_mcp_server:
            self._add_router(app, get_home_router(self))

        self._add_router(app, get_health_router(health_endpoint="/health"))
        self._add_router(app, get_base_router(self, settings=self.settings))
        self._add_router(app, get_agent_router(self, settings=self.settings, registry=self.registry))
        self._add_router(app, get_team_router(self, settings=self.settings, registry=self.registry))
        self._add_router(app, get_workflow_router(self, settings=self.settings))
        self._add_router(app, get_websocket_router(self, settings=self.settings))

        # Add A2A interface if relevant
        has_a2a_interface = False
        for interface in self.interfaces:
            if not has_a2a_interface and interface.__class__.__name__ == "A2A":
                has_a2a_interface = True
            interface_router = interface.get_router()
            self._add_router(app, interface_router)
        if self.a2a_interface and not has_a2a_interface:
            from agno.os.interfaces.a2a import A2A

            a2a_interface = A2A(agents=self.agents, teams=self.teams, workflows=self.workflows)
            self.interfaces.append(a2a_interface)
            self._add_router(app, a2a_interface.get_router())

    def _raise_if_duplicate_ids(self) -> None:
        """Check for duplicate IDs within each entity type.

        Raises:
            ValueError: If duplicate IDs are found within the same entity type
        """
        duplicate_ids: List[str] = []

        for entities in [self.agents, self.teams, self.workflows]:
            if not entities:
                continue
            seen_ids: set[str] = set()
            for entity in entities:
                entity_id = entity.id
                if entity_id is None:
                    continue
                if entity_id in seen_ids:
                    if entity_id not in duplicate_ids:
                        duplicate_ids.append(entity_id)
                else:
                    seen_ids.add(entity_id)

        if duplicate_ids:
            raise ValueError(f"Duplicate IDs found in AgentOS: {', '.join(repr(id_) for id_ in duplicate_ids)}")

    def _make_app(self, lifespan: Optional[Any] = None) -> FastAPI:
        return FastAPI(
            title=self.name or "Agno AgentOS",
            version=self.version or "1.0.0",
            description=self.description or "Your multi-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=lifespan,
        )

    def _initialize_agents(self) -> None:
        """Initialize and configure all agents for AgentOS usage."""
        if not self.agents:
            return
        for agent in self.agents:
            if isinstance(agent, RemoteAgent):
                continue
            # Set the default db to agents without their own
            if self.db is not None and agent.db is None:
                agent.db = self.db
            # Track all MCP tools to later handle their connection
            if agent.tools and isinstance(agent.tools, list):
                for tool in agent.tools:
                    # Checking if the tool is an instance of MCPTools, MultiMCPTools, or a subclass of those
                    if hasattr(type(tool), "__mro__"):
                        mro_names = {cls.__name__ for cls in type(tool).__mro__}
                        if mro_names & {"MCPTools", "MultiMCPTools"}:
                            if tool not in self.mcp_tools:
                                self.mcp_tools.append(tool)

            agent.initialize_agent()

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

            # Propagate run_hooks_in_background setting from AgentOS to agents
            agent._run_hooks_in_background = self.run_hooks_in_background

    def _initialize_teams(self) -> None:
        """Initialize and configure all teams for AgentOS usage."""
        if not self.teams:
            return

        for team in self.teams:
            if isinstance(team, RemoteTeam):
                continue

            # Set the default db to teams without their own
            if self.db is not None and team.db is None:
                team.db = self.db

            # Track all MCP tools recursively
            collect_mcp_tools_from_team(team, self.mcp_tools)

            team.initialize_team()

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

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

            # Propagate run_hooks_in_background setting to team and all nested members
            team.propagate_run_hooks_in_background(self.run_hooks_in_background)

    def _initialize_workflows(self) -> None:
        """Initialize and configure all workflows for AgentOS usage."""
        if not self.workflows:
            return

        for workflow in self.workflows:
            if isinstance(workflow, RemoteWorkflow):
                continue
            # Set the default db to workflows without their own
            if self.db is not None and workflow.db is None:
                workflow.db = self.db

            # Track MCP tools recursively in workflow members
            collect_mcp_tools_from_workflow(workflow, self.mcp_tools)

            if not workflow.id:
                workflow.id = generate_id_from_name(workflow.name)

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

            # Propagate run_hooks_in_background setting to workflow and all its step agents/teams
            workflow.propagate_run_hooks_in_background(self.run_hooks_in_background)

    def _populate_registry(self) -> None:
        """Populate the registry with code-defined agents and teams.

        This ensures that workflows loaded from DB can rehydrate their steps
        using code-defined agents/teams via the registry.
        """
        if self.registry is None:
            self.registry = Registry()

        if self.agents:
            existing_agent_ids = {getattr(a, "id", None) for a in self.registry.agents}
            for agent in self.agents:
                agent_id = getattr(agent, "id", None)
                if not isinstance(agent, RemoteAgent) and agent_id is not None and agent_id not in existing_agent_ids:
                    self.registry.agents.append(agent)
                    existing_agent_ids.add(agent_id)

        if self.teams:
            existing_team_ids = {getattr(t, "id", None) for t in self.registry.teams}
            for team in self.teams:
                team_id = getattr(team, "id", None)
                if not isinstance(team, RemoteTeam) and team_id is not None and team_id not in existing_team_ids:
                    self.registry.teams.append(team)
                    existing_team_ids.add(team_id)

    def _setup_tracing(self) -> None:
        """Set up OpenTelemetry tracing for this AgentOS.

        Uses the AgentOS db if provided, otherwise falls back to the first available
        database from agents/teams/workflows.
        """
        # Use AgentOS db if explicitly provided
        if self.db is not None:
            setup_tracing_for_os(db=self.db)
            return

        # Fall back to finding the first available database
        db: Optional[Union[BaseDb, AsyncBaseDb, RemoteDb]] = None

        for agent in self.agents or []:
            if agent.db:
                db = agent.db
                break

        if db is None:
            for team in self.teams or []:
                if team.db:
                    db = team.db
                    break

        if db is None:
            for workflow in self.workflows or []:
                if workflow.db:
                    db = workflow.db
                    break

        if db is None:
            log_warning(
                "tracing=True but no database found. "
                "Provide 'db' parameter to AgentOS or to at least one agent/team/workflow."
            )
            return

        setup_tracing_for_os(db=db)

    def get_app(self) -> FastAPI:
        if self.base_app:
            fastapi_app = self.base_app

            # Initialize MCP server if enabled
            if self.enable_mcp_server:
                from agno.os.mcp import get_mcp_server

                self._mcp_app = get_mcp_server(self)

            # Collect all lifespans that need to be combined
            lifespans = []

            # The user provided lifespan
            if self.lifespan:
                # Wrap the user lifespan with agent_os parameter
                wrapped_lifespan = self._add_agent_os_to_lifespan_function(self.lifespan)
                lifespans.append(wrapped_lifespan)

            # The provided app's existing lifespan
            if fastapi_app.router.lifespan_context:
                lifespans.append(fastapi_app.router.lifespan_context)

            # The MCP tools lifespan
            if self.mcp_tools:
                lifespans.append(partial(mcp_lifespan, mcp_tools=self.mcp_tools))

            # The /mcp server lifespan
            if self.enable_mcp_server and self._mcp_app:
                lifespans.append(self._mcp_app.lifespan)

            # The async database lifespan
            lifespans.append(partial(db_lifespan, agent_os=self))

            # The scheduler lifespan (after db so tables exist)
            if self._scheduler_enabled and self.db is not None:
                lifespans.append(partial(scheduler_lifespan, agent_os=self))

            # The httpx client cleanup lifespan (should be last to close after other lifespans)
            lifespans.append(http_client_lifespan)

            # Combine lifespans and set them in the app
            if lifespans:
                fastapi_app.router.lifespan_context = _combine_app_lifespans(lifespans)

        else:
            lifespans = []

            # User provided lifespan
            if self.lifespan:
                lifespans.append(self._add_agent_os_to_lifespan_function(self.lifespan))

            # MCP tools lifespan
            if self.mcp_tools:
                lifespans.append(partial(mcp_lifespan, mcp_tools=self.mcp_tools))

            # MCP server lifespan
            if self.enable_mcp_server:
                from agno.os.mcp import get_mcp_server

                self._mcp_app = get_mcp_server(self)
                lifespans.append(self._mcp_app.lifespan)

            # Async database initialization lifespan
            lifespans.append(partial(db_lifespan, agent_os=self))  # type: ignore

            # The scheduler lifespan (after db so tables exist)
            if self._scheduler_enabled and self.db is not None:
                lifespans.append(partial(scheduler_lifespan, agent_os=self))

            # The httpx client cleanup lifespan (should be last to close after other lifespans)
            lifespans.append(http_client_lifespan)

            final_lifespan = _combine_app_lifespans(lifespans) if lifespans else None
            fastapi_app = self._make_app(lifespan=final_lifespan)

        self._add_built_in_routes(app=fastapi_app)

        self._auto_discover_databases()
        self._auto_discover_knowledge_instances()

        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),
            get_traces_router(dbs=self.dbs),
            get_database_router(self, settings=self.settings),
        ]
        # Routes that require a database
        if self.db is not None:
            if isinstance(self.db, BaseDb):
                routers.append(get_components_router(os_db=self.db, registry=self.registry))
            else:
                routers.append(_get_disabled_feature_router("/components", "Components", "sync db (BaseDb)"))
            routers.append(get_schedule_router(os_db=self.db, settings=self.settings))
            routers.append(get_approval_router(os_db=self.db, settings=self.settings))
        else:
            log_debug("Components, Scheduler, and Approval routers not enabled: requires a db to be provided to AgentOS")
            for prefix, tag in [("/components", "Components"), ("/schedules", "Schedules"), ("/approvals", "Approvals")]:
                routers.append(_get_disabled_feature_router(prefix, tag, "db"))

        # Registry router
        if self.registry is not None:
            routers.append(get_registry_router(registry=self.registry))
        else:
            log_debug("Registry router not enabled: requires a registry to be provided to AgentOS")
            routers.append(_get_disabled_feature_router("/registry", "Registry", "registry"))

        for router in routers:
            self._add_router(fastapi_app, router)

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

        if not self._app_set:

            @fastapi_app.exception_handler(RequestValidationError)
            async def validation_exception_handler(_: Request, exc: RequestValidationError) -> JSONResponse:
                log_error(f"Validation error (422): {exc.errors()}")
                return JSONResponse(
                    status_code=422,
                    content={"detail": exc.errors()},
                )

            @fastapi_app.exception_handler(HTTPException)
            async def http_exception_handler(_, exc: HTTPException) -> JSONResponse:
                log_error(f"HTTP exception: {exc.status_code} {exc.detail}")
                return JSONResponse(
                    status_code=exc.status_code,
                    content={"detail": str(exc.detail)},
                )

            @fastapi_app.exception_handler(HTTPStatusError)
            async def http_status_error_handler(_: Request, exc: HTTPStatusError) -> JSONResponse:
                status_code = exc.response.status_code
                detail = exc.response.text
                log_error(f"Downstream server returned HTTP status error: {status_code} {detail}")
                return JSONResponse(
                    status_code=status_code,
                    content={"detail": detail},
                )

            @fastapi_app.exception_handler(Exception)
            async def general_exception_handler(_: Request, exc: Exception) -> JSONResponse:
                import traceback

                log_error(f"Unhandled exception:\n{traceback.format_exc(limit=5)}")

                return JSONResponse(
                    status_code=getattr(exc, "status_code", 500),
                    content={"detail": str(exc)},
                )

        # Update CORS middleware
        update_cors_middleware(fastapi_app, self.cors_allowed_origins)  # type: ignore

        # Set agent_os_id and cors_allowed_origins on app state
        # This allows middleware (like JWT) to access these values
        fastapi_app.state.agent_os_id = self.id
        fastapi_app.state.cors_allowed_origins = self.cors_allowed_origins

        # Store internal service token for scheduler auth bypass
        if self._internal_service_token:
            fastapi_app.state.internal_service_token = self._internal_service_token

        # Add JWT middleware if authorization is enabled
        if self.authorization:
            # Set authorization_enabled flag on settings so security key validation is skipped
            self.settings.authorization_enabled = True

            jwt_configured = bool(getenv("JWT_VERIFICATION_KEY") or getenv("JWT_JWKS_FILE"))
            security_key_set = bool(self.settings.os_security_key)
            if jwt_configured and security_key_set:
                log_warning(
                    "Both JWT configuration (JWT_VERIFICATION_KEY or JWT_JWKS_FILE) and OS_SECURITY_KEY are set. "
                    "With authorization=True, only JWT authorization will be used. "
                    "Consider removing OS_SECURITY_KEY from your environment."
                )

            self._add_jwt_middleware(fastapi_app)

        # Add trailing slash normalization middleware
        from agno.os.middleware.trailing_slash import TrailingSlashMiddleware

        fastapi_app.add_middleware(TrailingSlashMiddleware)

        return fastapi_app

    def _add_jwt_middleware(self, fastapi_app: FastAPI) -> None:
        from agno.os.middleware.jwt import JWTMiddleware, JWTValidator

        verify_audience = False
        jwks_file = None
        verification_keys = None
        algorithm = "RS256"

        if self.authorization_config:
            algorithm = self.authorization_config.algorithm or "RS256"
            verification_keys = self.authorization_config.verification_keys
            jwks_file = self.authorization_config.jwks_file
            verify_audience = self.authorization_config.verify_audience or False

        log_info(f"Adding JWT middleware for authorization (algorithm: {algorithm})")

        # Create validator and store on app.state for WebSocket access
        jwt_validator = JWTValidator(
            verification_keys=verification_keys,
            jwks_file=jwks_file,
            algorithm=algorithm,
        )
        fastapi_app.state.jwt_validator = jwt_validator

        # Add middleware to stack
        fastapi_app.add_middleware(
            JWTMiddleware,
            verification_keys=verification_keys,
            jwks_file=jwks_file,
            algorithm=algorithm,
            authorization=self.authorization,
            verify_audience=verify_audience,
        )

    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 _add_router(self, fastapi_app: FastAPI, router: APIRouter) -> None:
        """Add a router to the FastAPI app, avoiding route conflicts.

        Args:
            router: The APIRouter to add
        """

        conflicts = find_conflicting_routes(fastapi_app, router)
        conflicting_routes = [conflict["route"] for conflict in conflicts]

        if conflicts and self._app_set:
            if self.on_route_conflict == "preserve_base_app":
                # Skip conflicting AgentOS routes, prefer user's existing routes
                for conflict in conflicts:
                    methods_str = ", ".join(conflict["methods"])  # type: ignore
                    log_debug(
                        f"Skipping conflicting AgentOS route: {methods_str} {conflict['path']} - "
                        f"Using existing custom route instead"
                    )

                # Create a new router without the conflicting routes
                filtered_router = APIRouter()
                for route in router.routes:
                    if route not in conflicting_routes:
                        filtered_router.routes.append(route)

                # Use the filtered router if it has any routes left
                if filtered_router.routes:
                    fastapi_app.include_router(filtered_router)

            elif self.on_route_conflict == "preserve_agentos":
                # Log warnings but still add all routes (AgentOS routes will override)
                for conflict in conflicts:
                    methods_str = ", ".join(conflict["methods"])  # type: ignore
                    log_warning(
                        f"Route conflict detected: {methods_str} {conflict['path']} - "
                        f"AgentOS route will override existing custom route"
                    )

                # Remove conflicting routes
                for route in fastapi_app.routes:
                    for conflict in conflicts:
                        if isinstance(route, APIRoute):
                            if route.path == conflict["path"] and list(route.methods) == list(conflict["methods"]):  # type: ignore
                                fastapi_app.routes.pop(fastapi_app.routes.index(route))

                fastapi_app.include_router(router)

            elif self.on_route_conflict == "error":
                conflicting_paths = [conflict["path"] for conflict in conflicts]
                raise ValueError(f"Route conflict detected: {conflicting_paths}")

        else:
            # No conflicts, add router normally
            fastapi_app.include_router(router)

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

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

        dbs: Dict[str, List[Union[BaseDb, AsyncBaseDb, RemoteDb]]] = {}
        knowledge_dbs: Dict[
            str, List[Union[BaseDb, AsyncBaseDb, RemoteDb]]
        ] = {}  # Track databases specifically used for knowledge

        for agent in self.agents or []:
            if agent.db:
                self._register_db_with_validation(dbs, agent.db)
            agent_contents_db = getattr(agent.knowledge, "contents_db", None) if agent.knowledge else None
            if agent_contents_db:
                self._register_db_with_validation(knowledge_dbs, agent_contents_db)

        for team in self.teams or []:
            if team.db:
                self._register_db_with_validation(dbs, team.db)
            team_contents_db = getattr(team.knowledge, "contents_db", None) if team.knowledge else None
            if team_contents_db:
                self._register_db_with_validation(knowledge_dbs, team_contents_db)

        for workflow in self.workflows or []:
            if workflow.db:
                self._register_db_with_validation(dbs, workflow.db)

        for knowledge_base in self.knowledge or []:
            if knowledge_base.contents_db:
                self._register_db_with_validation(knowledge_dbs, knowledge_base.contents_db)

        for interface in self.interfaces or []:
            if interface.agent and interface.agent.db:
                self._register_db_with_validation(dbs, interface.agent.db)
            elif interface.team and interface.team.db:
                self._register_db_with_validation(dbs, interface.team.db)

        # Register AgentOS db if provided
        if self.db is not None:
            self._register_db_with_validation(dbs, self.db)

        self.dbs = dbs
        self.knowledge_dbs = knowledge_dbs

        # Initialize all discovered databases
        if self.auto_provision_dbs:
            self._pending_async_db_init = True

    def _initialize_sync_databases(self) -> None:
        """Initialize sync databases."""
        from itertools import chain

        unique_dbs = list(
            {
                id(db): db
                for db in chain(
                    chain.from_iterable(self.dbs.values()), chain.from_iterable(self.knowledge_dbs.values())
                )
            }.values()
        )

        for db in unique_dbs:
            if isinstance(db, AsyncBaseDb):
                continue  # Skip async dbs

            try:
                if hasattr(db, "_create_all_tables") and callable(db._create_all_tables):
                    db._create_all_tables()
            except Exception as e:
                log_warning(f"Failed to initialize {db.__class__.__name__} (id: {db.id}): {e}")

    async def _initialize_async_databases(self) -> None:
        """Initialize async databases."""

        from itertools import chain

        unique_dbs = list(
            {
                id(db): db
                for db in chain(
                    chain.from_iterable(self.dbs.values()), chain.from_iterable(self.knowledge_dbs.values())
                )
            }.values()
        )

        for db in unique_dbs:
            if not isinstance(db, AsyncBaseDb):
                continue  # Skip sync dbs

            try:
                if hasattr(db, "_create_all_tables") and callable(db._create_all_tables):
                    await db._create_all_tables()
            except Exception as e:
                log_warning(f"Failed to initialize async {db.__class__.__name__} (id: {db.id}): {e}")

    async def _close_databases(self) -> None:
        """Close all database connections and release connection pools."""
        from itertools import chain

        if not hasattr(self, "dbs") or not hasattr(self, "knowledge_dbs"):
            return

        unique_dbs = list(
            {
                id(db): db
                for db in chain(
                    chain.from_iterable(self.dbs.values()), chain.from_iterable(self.knowledge_dbs.values())
                )
            }.values()
        )

        for db in unique_dbs:
            try:
                if hasattr(db, "close") and callable(db.close):
                    if isinstance(db, AsyncBaseDb):
                        await db.close()
                    else:
                        db.close()
            except Exception as e:
                log_warning(f"Failed to close {db.__class__.__name__} (id: {db.id}): {e}")

    def _get_db_table_names(self, db: BaseDb) -> Dict[str, str]:
        """Get the table names for a database"""
        table_names = {
            "session_table_name": db.session_table_name,
            "culture_table_name": db.culture_table_name,
            "memory_table_name": db.memory_table_name,
            "metrics_table_name": db.metrics_table_name,
            "evals_table_name": db.eval_table_name,
            "knowledge_table_name": db.knowledge_table_name,
        }
        return {k: v for k, v in table_names.items() if v is not None}

    def _register_db_with_validation(
        self,
        registered_dbs: Dict[str, List[Union[BaseDb, AsyncBaseDb, RemoteDb]]],
        db: Union[BaseDb, AsyncBaseDb, RemoteDb],
    ) -> None:
        """Register a database in the contextual OS after validating it is not conflicting with registered databases"""
        if db.id in registered_dbs:
            registered_dbs[db.id].append(db)
        else:
            registered_dbs[db.id] = [db]

    def _auto_discover_knowledge_instances(self) -> None:
        """Auto-discover the knowledge instances used by all contextual agents, teams and workflows."""
        seen_instances: set[int] = set()  # Track by object identity
        knowledge_instances: List[Union[Knowledge, RemoteKnowledge]] = []

        def _add_knowledge_if_not_duplicate(knowledge: Any) -> None:
            """Add knowledge instance if it's not already in the list (by object identity)."""
            # Only handle Knowledge and RemoteKnowledge instances that have contents_db
            contents_db = getattr(knowledge, "contents_db", None)
            if not contents_db:
                return
            # Deduplicate by object identity to allow multiple knowledge instances with the same contents_db
            if id(knowledge) in seen_instances:
                return
            seen_instances.add(id(knowledge))
            # Only append if it's a Knowledge or RemoteKnowledge instance
            if isinstance(knowledge, (Knowledge, RemoteKnowledge)):
                knowledge_instances.append(knowledge)

        for agent in self.agents or []:
            if agent.knowledge:
                _add_knowledge_if_not_duplicate(agent.knowledge)

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

        for knowledge_base in self.knowledge or []:
            _add_knowledge_if_not_duplicate(knowledge_base)

        self.knowledge_instances = knowledge_instances

        # Validate that all knowledge instances have unique names
        # Duplicate names cause content isolation issues since linked_to uses the name
        self._validate_knowledge_instance_names()

    def _validate_knowledge_instance_names(self) -> None:
        """Validate that all knowledge instances have unique names.

        Raises:
            ValueError: If duplicate knowledge instance names are detected.
        """
        # Track seen combinations of (name, db_id, table) to detect true duplicates
        # Same name is OK if using different contents_db or different table
        seen_combinations: dict[tuple[str, str, str], str] = {}  # (name, db_id, table) -> description
        duplicates: list[str] = []

        for knowledge in self.knowledge_instances:
            contents_db = getattr(knowledge, "contents_db", None)
            if not contents_db:
                continue

            db_id = getattr(contents_db, "id", None)
            if not db_id:
                continue

            # Get the name (with fallback)
            knowledge_name = getattr(knowledge, "name", None) or f"knowledge_{db_id}"
            table_name = getattr(contents_db, "knowledge_table_name", "unknown")

            # Create unique key based on name + db + table
            key = (knowledge_name, db_id, table_name)

            if key in seen_combinations:
                duplicates.append(f"'{knowledge_name}' in table '{table_name}'")
            else:
                seen_combinations[key] = table_name

        if duplicates:
            unique_duplicates = list(set(duplicates))
            error_msg = (
                f"Duplicate knowledge instances detected:\n"
                f"  {', '.join(unique_duplicates)}\n\n"
                "Each knowledge instance must have a unique combination of (knowledgename, database, table). "
                "To fix this, give each knowledge instance a unique `name` parameter."
            )
            log_error(error_msg)
            raise ValueError(error_msg)

    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 = []

        dbs_with_specific_config = [db.db_id for db in session_config.dbs]
        for db_id, dbs in self.dbs.items():
            if db_id not in dbs_with_specific_config:
                # Collect unique table names from all databases with the same id
                unique_tables = list(set(db.session_table_name for db in dbs))
                session_config.dbs.append(
                    DatabaseConfig(
                        db_id=db_id,
                        domain_config=SessionDomainConfig(display_name=db_id),
                        tables=unique_tables,
                    )
                )

        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 = []

        dbs_with_specific_config = [db.db_id for db in memory_config.dbs]

        for db_id, dbs in self.dbs.items():
            if db_id not in dbs_with_specific_config:
                # Collect unique table names from all databases with the same id
                unique_tables = list(set(db.memory_table_name for db in dbs))
                memory_config.dbs.append(
                    DatabaseConfig(
                        db_id=db_id,
                        domain_config=MemoryDomainConfig(display_name=db_id),
                        tables=unique_tables,
                    )
                )

        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 = []
        if knowledge_config.knowledge_instances is None:
            knowledge_config.knowledge_instances = []

        # Track seen knowledge IDs to deduplicate
        seen_knowledge_ids: set[str] = set()

        # Build flat list of knowledge instances
        for knowledge in self.knowledge_instances:
            contents_db = getattr(knowledge, "contents_db", None)
            if not contents_db:
                continue

            db_id = getattr(contents_db, "id", None)
            if not db_id:
                continue

            table_name = getattr(contents_db, "knowledge_table_name", "unknown")
            knowledge_name = getattr(knowledge, "name", None) or f"knowledge_{db_id}"
            knowledge_id = _generate_knowledge_id(knowledge_name, db_id, table_name)

            # Skip if already processed (deduplicate by knowledge_id)
            if knowledge_id in seen_knowledge_ids:
                continue
            seen_knowledge_ids.add(knowledge_id)

            instance_config = KnowledgeInstanceConfig(
                id=knowledge_id,
                name=knowledge_name,
                description=getattr(knowledge, "description", None),
                db_id=db_id,
                table=table_name,
            )
            knowledge_config.knowledge_instances.append(instance_config)

        # Build KnowledgeDatabaseConfig for each db with its tables (as strings)
        dbs_with_specific_config = [db.db_id for db in knowledge_config.dbs]

        for db_id, dbs in self.knowledge_dbs.items():
            if db_id not in dbs_with_specific_config:
                # Get all unique table names for this db
                unique_tables = list(set(db.knowledge_table_name for db in dbs))

                knowledge_config.dbs.append(
                    KnowledgeDatabaseConfig(
                        db_id=db_id,
                        domain_config=KnowledgeDomainConfig(display_name=db_id),
                        tables=unique_tables,
                    )
                )

        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 = []

        dbs_with_specific_config = [db.db_id for db in metrics_config.dbs]

        for db_id, dbs in self.dbs.items():
            if db_id not in dbs_with_specific_config:
                # Collect unique table names from all databases with the same id
                unique_tables = list(set(db.metrics_table_name for db in dbs))
                metrics_config.dbs.append(
                    DatabaseConfig(
                        db_id=db_id,
                        domain_config=MetricsDomainConfig(display_name=db_id),
                        tables=unique_tables,
                    )
                )

        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 = []

        dbs_with_specific_config = [db.db_id for db in evals_config.dbs]

        for db_id, dbs in self.dbs.items():
            if db_id not in dbs_with_specific_config:
                # Collect unique table names from all databases with the same id
                unique_tables = list(set(db.eval_table_name for db in dbs))
                evals_config.dbs.append(
                    DatabaseConfig(
                        db_id=db_id,
                        domain_config=EvalsDomainConfig(display_name=db_id),
                        tables=unique_tables,
                    )
                )

        return evals_config

    def _get_traces_config(self) -> TracesConfig:
        traces_config = self.config.traces if self.config and self.config.traces else TracesConfig()

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

        dbs_with_specific_config = [db.db_id for db in traces_config.dbs]

        # If AgentOS db is explicitly set, only use that database for traces
        if self.db is not None:
            if self.db.id not in dbs_with_specific_config:
                traces_config.dbs.append(
                    DatabaseConfig(
                        db_id=self.db.id,
                        domain_config=TracesDomainConfig(display_name=self.db.id),
                    )
                )
        else:
            # Fall back to all discovered databases
            for db_id in self.dbs.keys():
                if db_id not in dbs_with_specific_config:
                    traces_config.dbs.append(
                        DatabaseConfig(
                            db_id=db_id,
                            domain_config=TracesDomainConfig(display_name=db_id),
                        )
                    )

        return traces_config

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

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

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

        panel_group = [
            Align.center(f"[bold cyan]{public_endpoint}[/bold cyan]"),
            Align.center(f"\n\n[bold dark_orange]OS running on:[/bold dark_orange] http://{host}:{port}"),
        ]
        if self.authorization:
            panel_group.append(
                Align.center("\n\n[bold chartreuse3]:lock: JWT Authorization Enabled[/bold chartreuse3]")
            )
        elif bool(self.settings.os_security_key):
            panel_group.append(Align.center("\n\n[bold chartreuse3]:lock: Security Key Enabled[/bold chartreuse3]"))

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

        # Adding *.yaml to reload_includes to reload the app when the yaml config file changes.
        if reload and reload_includes is not None:
            reload_includes = ["*.yaml", "*.yml"]

        uvicorn.run(
            app=app,
            host=host,
            port=port,
            reload=reload,
            reload_includes=reload_includes,
            reload_excludes=reload_excludes,
            workers=workers,
            access_log=access_log,
            lifespan="on",
            **kwargs,
        )
