"""FastAPI application for icarus serve.

Exposes a WebSocket endpoint at /ws and a health check at /health.
All agent communication flows through the WebSocket using the same
JSON protocol as the desktop sidecar.
"""
from __future__ import annotations

import asyncio
import logging
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware

from icarus_server.auth.db import init_server_db
from icarus_server.auth.tokens import load_private_key, verify_access_token
from icarus_server.auth.local import LocalAuthProvider
from icarus_server.auth.rate_limit import RateLimiter
from icarus_server.auth.routes import create_auth_router
from icarus_server.config import SERVER_DIR, ServerConfig
from icarus_server.engine_pool import EnginePool
from icarus_server.handler import MessageHandler
from icarus_server.users import provision_user_dir
from icarus_server.ws import ConnectionManager

logger = logging.getLogger(__name__)


def create_app(config: ServerConfig) -> FastAPI:
    """Build and return the FastAPI app wired to an EnginePool."""
    manager = ConnectionManager()

    # Auth initialization
    server_dir = SERVER_DIR
    db_path = server_dir / "server.db"
    key_path = server_dir / "jwt_signing.pem"
    server_dir.mkdir(parents=True, exist_ok=True)
    init_server_db(db_path)
    if not key_path.exists():
        logger.error("JWT signing key not found at %s. Run 'icarus setup' first.", key_path)
        raise SystemExit(f"Missing {key_path}. Run 'icarus setup' first.")
    jwt_key = load_private_key(key_path)
    auth_provider = LocalAuthProvider(db_path)
    rate_limiter = RateLimiter()

    @asynccontextmanager
    async def lifespan(app: FastAPI):
        pool = EnginePool(server_dir)
        pool.set_engine_initializer(_make_engine_initializer(manager))
        handler = MessageHandler(pool, manager)
        app.state.pool = pool
        app.state.handler = handler
        app.state.manager = manager

        # Auth state
        app.state.db_path = db_path
        app.state.jwt_key = jwt_key
        app.state.auth_provider = auth_provider
        app.state.rate_limiter = rate_limiter
        app.state.config = config
        app.state.server_dir = server_dir

        pool.start_sweep()
        yield
        await pool.shutdown()

    app = FastAPI(lifespan=lifespan)
    # CORS: the auth flow uses HTTP fetch (discover, login, refresh) which
    # is subject to CORS in the Tauri WebView during dev (localhost:3000).
    # Production Tauri builds use tauri:// and bypass CORS entirely.
    # JWT auth handles security — origin restrictions add nothing.
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_methods=["*"],
        allow_headers=["*"],
    )
    app.include_router(create_auth_router())

    @app.get("/health")
    async def health():
        return {"status": "ok"}

    @app.websocket("/ws")
    async def websocket_endpoint(ws: WebSocket):
        await ws.accept()

        # Auth handshake: first message must be {"type": "auth", "token": "<JWT>"}
        try:
            first = await asyncio.wait_for(ws.receive_json(), timeout=10)
        except (asyncio.TimeoutError, WebSocketDisconnect):
            await ws.close(code=1008)
            return

        if first.get("type") != "auth":
            await ws.send_json({"type": "auth_error", "message": "Invalid auth message"})
            await ws.close(code=1008)
            return

        # Verify JWT token
        payload = verify_access_token(jwt_key, first.get("token"))
        if not payload:
            await ws.send_json({"type": "auth_error", "message": "Invalid or expired token"})
            await ws.close(code=1008)
            return

        user_id = payload["sub"]

        # Look up user role for authorization
        from icarus_server.auth.db import get_user_by_id
        user = await asyncio.to_thread(get_user_by_id, db_path, user_id)
        if not user:
            await ws.send_json({"type": "auth_error", "message": "User not found"})
            await ws.close(code=1008)
            return
        user_role = user.role

        # Provision user dir on first connection
        await asyncio.to_thread(provision_user_dir, server_dir, user_id)

        await ws.send_json({"type": "auth_ok"})

        accepted = await manager.accept(ws, user_id, role=user_role)
        if not accepted:
            await ws.send_json({"type": "error", "message": "Connection limit reached"})
            await ws.close(code=1013)
            return

        pool: EnginePool = app.state.pool
        handler: MessageHandler = app.state.handler
        try:
            # Lazy-create engine and bump refcount for this connection
            await pool.get_or_create(user_id)
            while True:
                data = await ws.receive_json()
                await handler.handle(ws, data)
        except WebSocketDisconnect:
            pass
        except Exception:
            logger.exception("WebSocket error")
        finally:
            await pool.release(user_id)
            manager.disconnect(ws)

    return app


# ── Per-engine initializer (called by EnginePool on boot) ───

def _make_engine_initializer(manager: ConnectionManager):
    """Return a callback that installs MCP + state-change callbacks on a new engine."""
    def init_engine(engine: Any, user_id: str) -> None:
        _install_mcp_callbacks(engine, manager, user_id=user_id)

        def _on_thread_state_change(tid, _old, new):
            owner = manager.get_thread_owner(tid)
            data = {"type": "thread_state_changed", "thread_id": tid, "state": new.value}
            if owner:
                asyncio.ensure_future(manager.broadcast_to_user(owner, data))

        engine.state_manager.on_change(_on_thread_state_change)

    return init_engine


# ── MCP callback wiring ─────────────────────────────────────

def _install_mcp_callbacks(engine: Any, manager: ConnectionManager, *, user_id: str) -> None:
    """Patch engine MCP callbacks to broadcast WS events to the owning user only."""
    orig_failure = engine._on_mcp_server_failure
    orig_recovery = engine._on_mcp_server_recovery
    orig_reconnect = engine._on_mcp_server_reconnect
    orig_boot_failures = engine._on_mcp_boot_failures

    def on_failure(name: str, msg: str) -> None:
        orig_failure(name, msg)
        asyncio.ensure_future(manager.broadcast_to_user(user_id, {
            "type": "connector_warning", "name": name,
            "message": f"Connector '{name}' \u2014 tools failing: {msg}",
        }))

    def on_recovery(name: str) -> None:
        orig_recovery(name)
        asyncio.ensure_future(manager.broadcast_to_user(user_id, {
            "type": "connector_recovered", "name": name,
            "message": f"Connector '{name}' recovered",
        }))

    def on_reconnect(name: str) -> None:
        orig_reconnect(name)
        asyncio.ensure_future(manager.broadcast_to_user(user_id, {
            "type": "connector_warning", "name": name,
            "message": f"Connector '{name}' reconnecting...",
        }))
        asyncio.ensure_future(engine.reload_mcp())

    async def on_boot_failures(failed: list[str]) -> None:
        await orig_boot_failures(failed)
        await _auto_disable_failed_connectors(engine, manager, failed, user_id=user_id)

    engine._on_mcp_server_failure = on_failure
    engine._on_mcp_server_recovery = on_recovery
    engine._on_mcp_server_reconnect = on_reconnect
    engine._on_mcp_boot_failures = on_boot_failures


async def _auto_disable_failed_connectors(
    engine: Any, manager: ConnectionManager, failed: list[str], *, user_id: str,
) -> None:
    """Auto-disable connectors that failed during boot."""
    from icarus.mcp.config import read_mcp_config, toggle_server, write_mcp_config

    config_path = engine._mcp_config_path()
    config = await asyncio.to_thread(read_mcp_config, config_path)
    toggled = False
    for name in failed:
        if name in config.get("mcpServers", {}):
            toggle_server(config, name)
            toggled = True
            error = (
                engine.mcp_manager.server_errors.get(name, "not reachable")
                if engine.mcp_manager else "not reachable"
            )
            await manager.broadcast_to_user(user_id, {
                "type": "connector_warning", "name": name,
                "message": f"Connector '{name}' disabled \u2014 {error}",
            })
    if toggled:
        await asyncio.to_thread(write_mcp_config, config_path, config)
