"""FastMCP server with async lifespan managing Postgres + Redis."""

from __future__ import annotations

import asyncio
import json
import logging
import os
import signal
import sys
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass

import asyncpg
import redis.asyncio as aioredis
from fastmcp import FastMCP

from loom.cloud.connections import create_pool, create_redis
from loom.cloud.secrets import maybe_load_gcp_secrets
from loom.config import LoomConfig, load_config
from loom.db.migrations import run_migrations
from loom.graph import cache
from loom.logging import setup_logging

logger = logging.getLogger(__name__)


@dataclass
class AppContext:
    pool: asyncpg.Pool
    redis: aioredis.Redis
    project_id: str
    config: LoomConfig = None
    project_dir: str = ""


async def _health_handler(
    reader: asyncio.StreamReader,
    writer: asyncio.StreamWriter,
    pool: asyncpg.Pool,
    redis_client: aioredis.Redis,
    project_id: str = "",
    config=None,
) -> None:
    """Handle HTTP requests for /health, /metrics, /dashboard, and /static."""
    try:
        data = await asyncio.wait_for(reader.read(8192), timeout=5.0)

        # Parse request path
        request_line = data.decode("utf-8", errors="replace").split("\r\n")[0]
        path = request_line.split(" ")[1] if " " in request_line else "/health"

        # Dashboard and static routes
        if path.startswith("/dashboard") or path.startswith("/static/"):
            from loom.dashboard.app import handle_dashboard_request
            await handle_dashboard_request(
                writer, pool, redis_client, project_id, config, path,
            )
            return

        if path == "/metrics":
            from loom.observability.metrics import metrics_handler
            body = metrics_handler()
            response = (
                "HTTP/1.1 200 OK\r\n"
                "Content-Type: text/plain; charset=utf-8\r\n"
                f"Content-Length: {len(body)}\r\n"
                "\r\n"
                f"{body}"
            )
        else:
            # Health check (existing logic)
            pg_ok = False
            try:
                async with pool.acquire() as conn:
                    await conn.fetchval("SELECT 1")
                pg_ok = True
            except Exception:
                pass

            redis_ok = False
            try:
                await redis_client.ping()
                redis_ok = True
            except Exception:
                pass

            healthy = pg_ok and redis_ok
            body = json.dumps({
                "status": "healthy" if healthy else "unhealthy",
                "postgres": "ok" if pg_ok else "error",
                "redis": "ok" if redis_ok else "error",
            })
            status = "200 OK" if healthy else "503 Service Unavailable"
            response = (
                f"HTTP/1.1 {status}\r\n"
                "Content-Type: application/json\r\n"
                f"Content-Length: {len(body)}\r\n"
                "\r\n"
                f"{body}"
            )

        writer.write(response.encode())
        await writer.drain()
    except Exception:
        pass
    finally:
        writer.close()


@asynccontextmanager
async def lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
    """Initialize pool, run migrations, connect Redis, rebuild cache."""
    # 1. Load GCP secrets before config (sets env vars)
    maybe_load_gcp_secrets()

    # 2. Load config
    # LOOM_PROJECT_DIR sets the codebase root — used to find .loom/config.yaml
    # and as the scan root for decomposition. Defaults to cwd if unset.
    project_dir = os.environ.get("LOOM_PROJECT_DIR", os.getcwd())
    config = load_config(project_dir)

    # 3. Setup structured logging
    setup_logging(config.log_level)
    logger.info(
        "Config loaded from %s (changes require MCP server restart). project_id=%s, api_key=%s",
        project_dir,
        config.project_id or "(none)",
        "set" if config.skills.api_key else "NOT SET",
    )

    # 3b. Setup tracing
    from loom.observability.tracing import setup_tracing
    setup_tracing(service_name="loom")

    # 4. Postgres via factory
    pool = await create_pool(config.database)
    await run_migrations(pool)

    # 5. Redis via factory
    redis_conn = await create_redis(config.redis)

    # 6. Rebuild cache
    if config.project_id:
        await cache.rebuild_cache(redis_conn, pool, config.project_id)

    # 7. Start health check server
    health_server = None
    try:
        health_server = await asyncio.start_server(
            lambda r, w: _health_handler(
                r, w, pool, redis_conn, config.project_id, config,
            ),
            host="0.0.0.0",
            port=config.mcp.health_port,
        )
        logger.info("Health check listening on port %d", config.mcp.health_port)
    except OSError:
        logger.warning(
            "Could not start health check on port %d (port in use?)",
            config.mcp.health_port,
        )

    # 8. Register SIGTERM handler for graceful shutdown
    shutdown_event = asyncio.Event()

    def _sigterm_handler() -> None:
        logger.info("Received SIGTERM, shutting down gracefully...")
        shutdown_event.set()

    loop = asyncio.get_running_loop()
    if sys.platform != "win32":
        loop.add_signal_handler(signal.SIGTERM, _sigterm_handler)

    try:
        yield AppContext(
            pool=pool, redis=redis_conn, project_id=config.project_id,
            config=config, project_dir=project_dir,
        )
    finally:
        logger.info("Shutting down MCP server...")
        if health_server is not None:
            health_server.close()
            await health_server.wait_closed()
        await redis_conn.aclose()
        await pool.close()
        logger.info("Shutdown complete")


mcp = FastMCP(name="loom", lifespan=lifespan)

# Import tools to register them on the server
import loom.mcp.tools  # noqa: E402, F401
