"""Process entrypoint.

Phase 1.4 — async ``main``: install uvloop, load the user config, build
registries + per-agent backends (only ``RaycastBackend`` so far), wire
each frontend with a ``GatewayRuntime``, and run all
``frontend.serve()`` coroutines concurrently. Without any frontends we
still print the Phase 0 DoD line and exit cleanly so the bare skeleton
keeps working.

Phase 2.1 — when the user declares any ``McpServer``, we additionally
build the internal MCP aggregator app and run it on
``127.0.0.1:INTERNAL_MCP_PORT`` as another task inside the same
TaskGroup. URLs are surfaced through ``GatewayRuntime.mcp_internal_urls``
so Phase 2.2's ClaudeCode adapter can find them.

Phase 3 — the same aggregator backs the external ``McpServerFrontend``;
``cli`` doesn't have to know that the frontend reverse-proxies into it,
it just keeps the aggregator running for anyone who needs it.
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
import signal
from contextlib import AsyncExitStack
from typing import TYPE_CHECKING

import uvicorn
import uvloop
from dotenv import load_dotenv
from raycast_api import Client as RaycastClient
from raycast_api.config import Config as RaycastConfig

from beaver_gateway import config_loader
from beaver_gateway.agents.claude import ClaudeAgent
from beaver_gateway.agents.raycast import RaycastAgent
from beaver_gateway.backends.claude_sdk import (
    ClaudeSdkBackend,
    RunnerConfig,
    UsageEvent,
)
from beaver_gateway.backends.raycast import RaycastBackend
from beaver_gateway.core.auth import TokenStore
from beaver_gateway.core.registry import AgentRegistry, McpRegistry
from beaver_gateway.frontends.base import GatewayRuntime
from beaver_gateway.mcp.internal_app import build_internal_app
from beaver_gateway.settings import Settings
from beaver_gateway.storage import Database, PostgresSessionStore, Usage, append_usage

if TYPE_CHECKING:
    from fastmcp import FastMCP
    from fastmcp.tools.base import Tool as FastMCPTool
    from starlette.applications import Starlette

    from beaver_gateway.backends.base import Backend
    from beaver_gateway.mcp.types import McpServerT


_log = logging.getLogger("beaver_gateway.cli")


def main() -> None:
    """Sync wrapper: uvloop loop factory + asyncio.run."""
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
    )
    _install_sigterm_handler()
    asyncio.run(_async_main(), loop_factory=uvloop.new_event_loop)


def _install_sigterm_handler() -> None:
    """Turn SIGTERM into a normal interpreter exit.

    Python's default SIGTERM disposition kills the process outright, so
    neither ``AsyncExitStack`` unwinding nor ``atexit`` hooks run. That
    matters because every ``claude`` we spawn lives in its own session
    (ptyprocess calls ``setsid``), which makes it immune to the signal
    that took us down — a hard SIGTERM leaves one orphaned CLI per live
    session, each holding hundreds of MB. Raising ``KeyboardInterrupt``
    instead routes ``docker stop`` / ``systemctl stop`` through the same
    shutdown path as Ctrl-C, which does reap them.

    Only installed when we own the main thread's signal handlers; under
    an embedding host that isn't ours to take.
    """

    def _raise_interrupt(_signum: int, _frame: object) -> None:
        raise KeyboardInterrupt

    with contextlib.suppress(ValueError, OSError):
        signal.signal(signal.SIGTERM, _raise_interrupt)


async def _async_main() -> None:
    # Populate ``os.environ`` from ``.env`` before anything else so the
    # user's ``config.py`` can read its own secrets via ``os.environ[...]``
    # (Firefly PAT, third-party MCP creds, etc.). ``Settings`` already
    # reads ``.env`` independently via pydantic-settings, but that path
    # populates Settings fields, not the process environment.
    # ``override=False``: real env vars (Docker, systemd) win over .env.
    load_dotenv(override=False)

    settings = Settings()  # ty: ignore[missing-argument]
    gateway = config_loader.load(settings.config_path)

    agents = AgentRegistry(gateway.agents)
    mcps = McpRegistry(gateway.mcps)

    # Phase 4.1 — open the async DB and run create_all once. Engine
    # pool is process-wide; ``dispose()`` after the TaskGroup unwinds.
    db = Database(settings.database_url)
    await db.create_all()

    # Phase 4.2 — TokenStore now reads from the DB (in-memory cache
    # primed at start, TTL-refreshed, last_used_at flushed by a
    # background task). BOOTSTRAP_TOKENS layers on top so first-run /
    # examples still work without DB writes.
    token_store = TokenStore(
        db, bootstrap=TokenStore.parse_bootstrap(settings.bootstrap_tokens)
    )

    async with AsyncExitStack() as stack:
        stack.push_async_callback(db.dispose)
        await token_store.start()
        stack.push_async_callback(token_store.stop)
        # Internal MCP URLs must exist before we construct any
        # ClaudeSdkBackend - adapters bake the URLs into their
        # ``mcp_servers`` at construction time. The
        # ``mcp_servers`` map is used by the Raycast backend, which
        # needs in-process ``list_tools`` / ``call_tool`` access (the
        # Raycast wire has no native MCP concept).
        internal_app, internal_urls, mcp_servers = _build_internal_mcp(
            gateway.mcps, settings=settings
        )

        # Prefetch tool catalogs for every MCP so RaycastAgent requests
        # don't pay a per-turn list_tools roundtrip and so a broken MCP
        # surfaces at startup instead of mid-conversation.
        mcp_tools = await _prefetch_mcp_tools(mcp_servers)

        backends: dict[str, Backend] = await _build_backends(
            settings=settings,
            agents=agents,
            stack=stack,
            db=db,
            mcp_internal_urls=internal_urls,
            mcp_servers=mcp_servers,
            mcp_tools=mcp_tools,
        )

        runtime = GatewayRuntime(
            agents=agents,
            mcps=mcps,
            backends=backends,
            token_store=token_store,
            db=db,
            mcp_internal_urls=internal_urls,
            admin_user=settings.admin_user,
            admin_pass=settings.admin_pass,
            session_secret=settings.session_secret,
            frontends=tuple(gateway.frontends),
        )

        for fe in gateway.frontends:
            fe.configure(runtime)

        _log.info(
            "beaver-gateway: loaded %d agents, %d mcps, %d frontends",
            len(agents),
            len(mcps),
            len(gateway.frontends),
        )

        # Keep the Phase 0 DoD line on stdout for grep-friendly smoke
        # tests, in addition to the structured log line above.
        print(
            f"beaver-gateway: loaded {len(agents)} agents, "
            f"{len(mcps)} mcps, {len(gateway.frontends)} frontends"
        )

        if not gateway.frontends:
            # No external listeners → nothing to serve. The internal
            # MCP app has no consumer on its own, so we skip running
            # it in this path and exit cleanly (Phase 0 DoD).
            return

        async with asyncio.TaskGroup() as tg:
            if internal_app is not None:
                tg.create_task(_serve_internal_mcp(internal_app, settings=settings))
            for fe in gateway.frontends:
                tg.create_task(fe.serve())


def _build_internal_mcp(
    mcps: list[McpServerT], *, settings: Settings
) -> tuple[Starlette | None, dict[str, str], dict[str, FastMCP]]:
    """Build the aggregator app + URL map + server map, or empty equivalents.

    The URL map is always handed out (frontends may still introspect
    ``runtime.mcp_internal_urls`` even if nothing is configured); the
    app is ``None`` when there are no MCPs to mount, so the caller
    skips the uvicorn task entirely. The server map is the in-process
    handle the Raycast backend needs to splice MCP tools into its
    requests — empty when no MCPs are configured.
    """
    if not mcps:
        return None, {}, {}
    return build_internal_app(mcps, host="127.0.0.1", port=settings.internal_mcp_port)


async def _prefetch_mcp_tools(
    servers: dict[str, FastMCP],
) -> dict[str, list[FastMCPTool]]:
    """Eagerly enumerate tools per MCP so the Raycast loop has a static catalog.

    Each underlying proxy is allowed to fail independently — a broken
    MCP shouldn't take down the whole gateway. The result has one entry
    per MCP that responded; agents that ``expose_mcps`` a missing entry
    will simply expose no tools from it (logged once per request).
    """
    out: dict[str, list[FastMCPTool]] = {}
    for name, server in servers.items():
        try:
            out[name] = list(await server.list_tools())
        except Exception:  # noqa: BLE001 — proxy can raise any transport error; we degrade per-MCP rather than fail the whole gateway
            _log.exception("failed to list tools for MCP %r — skipping", name)
            out[name] = []
    return out


async def _serve_internal_mcp(app: Starlette, *, settings: Settings) -> None:
    """Run the internal MCP aggregator on loopback.

    Bound to ``127.0.0.1`` (never EXPOSE'd) — only the in-process
    ClaudeCode subprocess reaches it. Logged at ``warning`` level so
    we don't drown the gateway's own logs in per-request noise.
    """
    config = uvicorn.Config(
        app,
        host="127.0.0.1",
        port=settings.internal_mcp_port,
        log_level="warning",
        loop="uvloop",
    )
    server = uvicorn.Server(config)
    _log.info(
        "internal MCP aggregator on http://127.0.0.1:%d/mcp/<name>",
        settings.internal_mcp_port,
    )
    await server.serve()


async def _build_backends(
    *,
    settings: Settings,
    agents: AgentRegistry,
    stack: AsyncExitStack,
    db: Database,
    mcp_internal_urls: dict[str, str],
    mcp_servers: dict[str, FastMCP],
    mcp_tools: dict[str, list[FastMCPTool]],
) -> dict[str, Backend]:
    """Construct one backend per agent name.

    The Raycast ``Client`` is shared across every ``RaycastAgent``
    (bearer + device-id are process-wide), so we open it lazily - only
    when at least one ``RaycastAgent`` is present - and close it via
    the caller's exit stack.

    Each :class:`ClaudeAgent` gets its own :class:`ClaudeSdkBackend`
    (own live-session pool, own prompt and MCP set); all of them share
    the session store and the usage sink.
    """
    backends: dict[str, Backend] = {}

    raycast_agents = [a for a in agents if isinstance(a, RaycastAgent)]
    if raycast_agents:
        client = await _try_open_raycast_client(settings, stack)
        if client is not None:
            raycast_backend = RaycastBackend(
                client, mcp_servers=mcp_servers, mcp_tools=mcp_tools
            )
            for a in raycast_agents:
                backends[a.name] = raycast_backend

    session_store = PostgresSessionStore(db)
    runner = RunnerConfig(user=settings.claude_runner_user, home=settings.claude_home)
    mcp_tool_names = {
        name: [t.name for t in tools] for name, tools in mcp_tools.items()
    }

    async def record_usage(event: UsageEvent) -> None:
        row = Usage(
            agent_name=event.agent_name,
            conversation_id=event.conversation_id,
            session_id=event.session_id,
            model=event.model,
            effort=event.effort,
            input_tokens=event.usage.input_tokens,
            output_tokens=event.usage.output_tokens,
            cache_read_tokens=event.usage.cache_read_tokens,
            cache_creation_tokens=event.usage.cache_creation_tokens,
            cost_usd=event.usage.cost_usd,
            duration_ms=event.usage.duration_ms,
            num_turns=event.usage.num_turns,
        )
        try:
            async with db.session() as session:
                await append_usage(session, row)
        except Exception:  # noqa: BLE001
            _log.exception("usage write failed for %s", event.agent_name)

    for a in agents:
        if isinstance(a, ClaudeAgent):
            adapter = ClaudeSdkBackend(
                agent=a,
                mcp_internal_urls=mcp_internal_urls,
                session_store=session_store,
                mcp_tool_names=mcp_tool_names,
                runner=runner,
                usage_sink=record_usage,
            )
            await stack.enter_async_context(adapter)
            backends[a.name] = adapter

    return backends


async def _try_open_raycast_client(
    settings: Settings, stack: AsyncExitStack
) -> RaycastClient | None:
    """Open a ``raycast_api.Client`` if creds are available, else warn + skip.

    Skipping is gentler than failing startup: the user can bring up the
    gateway, list their agents, and still hit a ``ClaudeAgent``; affected
    ``RaycastAgent`` instances 503 with a specific message at request time.
    Missing ``raycast.json`` falls into the same bucket — first-run
    users won't have it yet.
    """
    if not settings.raycast_bearer:
        _log.warning(
            "RaycastAgent present but RAYCAST_BEARER is unset — those agents will 503"
        )
        return None
    if not settings.raycast_device_id:
        _log.warning(
            "RaycastAgent present but RAYCAST_DEVICE_ID is unset — those agents "
            "will 503 (generate once with `python -c 'import secrets; "
            "print(secrets.token_hex(32))'`)"
        )
        return None
    if not settings.raycast_config_path.exists():
        _log.warning(
            "RaycastAgent present but %s is missing — those agents will 503",
            settings.raycast_config_path,
        )
        return None

    config = RaycastConfig.load(settings.raycast_config_path)
    client = RaycastClient(
        config=config,
        bearer_token=settings.raycast_bearer,
        device_id=settings.raycast_device_id,
        locale=settings.raycast_locale,
    )
    return await stack.enter_async_context(client)