"""MCP server bootstrap — registers all tools, resources, and prompts."""

from __future__ import annotations

import asyncio
import hashlib
import logging
import os
import sys
import time

import httpx
from mcp.server import FastMCP
from mcp.server.fastmcp.server import StreamableHTTPASGIApp
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
from starlette.requests import Request
from starlette.responses import JSONResponse

from g8_mcp_server.auth import set_api_key
from g8_mcp_server.client import G8Client, G8ClientError
from g8_mcp_server.tools import campaigns, dev, forms, gtm, install, kb, platform, repos, snippet
from g8_mcp_server import resources, prompts

# Tool registration mode: "all" (default), "dev", or "gtm"
# - all:  registers every tool (backward compat)
# - dev:  repo-scoped tools for developers in Cursor/Windsurf/Claude Code
# - gtm:  org-scoped tools for campaign managers in Claude Desktop
MCP_MODE = os.getenv("G8_MCP_MODE", "all").lower()

logger = logging.getLogger(__name__)

# PropelAuth MCP token validation config
_PA_AUTH_URL = os.getenv("PROPELAUTH_AUTH_URL", "")
_PA_API_KEY = os.getenv("PROPELAUTH_API_KEY", "")
_INTROSPECT_URL = f"{_PA_AUTH_URL}/oauth/2.1/introspect"
_INTROSPECT_CLIENT_ID = os.getenv("MCP_INTROSPECTION_CLIENT_ID", "")
_INTROSPECT_CLIENT_SECRET = os.getenv("MCP_INTROSPECTION_CLIENT_SECRET", "")

# Redis cache prefix and TTL for JWT → API key mappings
_JWT_CACHE_PREFIX = "mcp:jwt_key:"
_JWT_CACHE_TTL = 14 * 86400  # 14 days (matches PropelAuth MCP session duration)


def _register_tools(server: FastMCP, client: G8Client, mode: str = "all") -> None:
    """Register tools based on mode: 'all', 'dev', or 'gtm'."""
    # Shared tools — available in every mode
    platform.register(server, client)

    if mode == "dev":
        dev.register(server, client)
        snippet.register(server, client)
        forms.register(server, client)
    elif mode == "gtm":
        gtm.register(server, client)
    else:
        # "all" — register legacy repo-scoped tools (backward compat)
        repos.register(server, client)
        install.register(server, client)
        campaigns.register(server, client)
        kb.register(server, client)

    resources.register(server, client)
    prompts.register(server)


def create_server() -> FastMCP:
    """Create MCP server for local stdio mode. Requires G8_API_KEY env var."""
    server = FastMCP("graph8")
    client = G8Client(require_key=True)
    _register_tools(server, client, mode=MCP_MODE)
    return server


# ---------------------------------------------------------------------------
# JWT → API key resolution (for PropelAuth MCP OAuth 2.1 tokens)
# ---------------------------------------------------------------------------

async def _introspect_token(token: str) -> dict | None:
    """Validate a token via PropelAuth's introspection endpoint."""
    async with httpx.AsyncClient(timeout=10) as client:
        resp = await client.post(
            _INTROSPECT_URL,
            data={"token": token},
            auth=(_INTROSPECT_CLIENT_ID, _INTROSPECT_CLIENT_SECRET),
        )
        if resp.status_code != 200:
            logger.warning("MCP introspection failed: %s", resp.status_code)
            return None
        data = resp.json()
        if not data.get("active"):
            return None
        return data


async def _fetch_user_org(user_id: str) -> str | None:
    """Fetch user's active org ID (PropelAuth UUID) from backend API."""
    async with httpx.AsyncClient(timeout=10) as client:
        resp = await client.get(
            f"{_PA_AUTH_URL}/api/backend/v1/user/{user_id}",
            headers={"Authorization": f"Bearer {_PA_API_KEY}"},
        )
        if resp.status_code != 200:
            logger.warning("MCP user fetch failed: %s", resp.status_code)
            return None
        user_data = resp.json()

    # Pick active org → first org fallback
    props = user_data.get("properties") or {}
    active_org_id = props.get("active_org_id")
    org_map = user_data.get("org_id_to_org_member_info") or {}

    if active_org_id:
        for pa_uuid, info in org_map.items():
            if info.get("legacy_org_id") == active_org_id:
                return pa_uuid
        if active_org_id in org_map:
            return active_org_id

    if org_map:
        return next(iter(org_map))
    return None


async def _create_api_key(org_id: str, user_id: str) -> str | None:
    """Create a temporary org API key via PropelAuth backend API."""
    expires_at = int(time.time()) + _JWT_CACHE_TTL
    async with httpx.AsyncClient(timeout=10) as client:
        resp = await client.post(
            f"{_PA_AUTH_URL}/api/backend/v1/end_user_api_keys",
            headers={
                "Authorization": f"Bearer {_PA_API_KEY}",
                "Content-Type": "application/json",
            },
            json={
                "org_id": org_id,
                "user_id": user_id,
                "expires_at_seconds": str(expires_at),
                "metadata": {"source": "mcp_oauth"},
            },
        )
        if resp.status_code not in (200, 201):
            logger.warning("MCP API key creation failed: %s %s", resp.status_code, resp.text)
            return None
        return resp.json().get("api_key_token")


def _redis():
    """Get Redis client (lazy import to avoid circular deps at module level)."""
    from shared import apps
    return apps.infra().redis()


async def _resolve_jwt_to_api_key(token: str) -> str | None:
    """Convert a PropelAuth MCP OAuth JWT to a cached API key.

    Flow (first request only — subsequent requests use Redis cache):
    1. Introspect token → validate + get user_id
    2. Fetch user details → resolve org
    3. Create a temporary org API key → cache in Redis
    """
    cache_key = f"{_JWT_CACHE_PREFIX}{hashlib.sha256(token.encode()).hexdigest()[:32]}"

    # Check cache
    try:
        redis = _redis()
        cached = await redis.get(cache_key)
        if cached:
            return cached.decode() if isinstance(cached, bytes) else cached
    except Exception:
        pass

    # Introspect
    intro = await _introspect_token(token)
    if not intro:
        return None

    user_id = intro.get("sub")
    if not user_id:
        return None

    # Resolve org
    org_id = await _fetch_user_org(user_id)
    if not org_id:
        logger.warning("MCP OAuth user %s has no organization", user_id)
        return None

    # Create API key
    api_key = await _create_api_key(org_id, user_id)
    if not api_key:
        return None

    # Cache for session duration
    try:
        redis = _redis()
        await redis.setex(cache_key, _JWT_CACHE_TTL, api_key)
    except Exception:
        pass

    logger.info("MCP OAuth: provisioned API key for user %s", intro.get("username", user_id))
    return api_key


# ---------------------------------------------------------------------------
# ASGI auth middleware
# ---------------------------------------------------------------------------

class _RemoteMCPApp:
    """ASGI app that serves MCP over HTTP with bearer auth.

    Combines authentication (API key + JWT/OAuth) with lazy lifecycle
    management for the MCP session manager.  The session manager is started
    as a background task on the first request — this works around FastAPI's
    ``app.mount()`` not forwarding lifespan events to sub-applications.
    """

    def __init__(self, handler: StreamableHTTPASGIApp, session_manager: StreamableHTTPSessionManager):
        self._handler = handler
        self._sm = session_manager
        self._started = False
        self._lock = asyncio.Lock()

    async def _ensure_started(self) -> None:
        """Start the session manager's task group once (lazy, thread-safe)."""
        if self._started:
            return
        async with self._lock:
            if self._started:
                return

            ready = asyncio.Event()

            async def _lifecycle() -> None:
                async with self._sm.run():
                    ready.set()
                    await asyncio.Event().wait()  # block until cancelled

            asyncio.create_task(_lifecycle())
            await ready.wait()
            self._started = True
            logger.info("MCP session manager started (lazy init)")

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http":
            return

        request = Request(scope)
        auth_header = request.headers.get("authorization", "")

        # --- Auth: require Bearer token ---
        if not auth_header.lower().startswith("bearer "):
            scheme = request.headers.get("x-forwarded-proto", request.url.scheme)
            backend_url = os.getenv("BACKEND_URL", f"{scheme}://{request.url.netloc}")
            resource_metadata_url = f"{backend_url}/.well-known/oauth-protected-resource"
            response = JSONResponse(
                {"error": "Missing or invalid Authorization header"},
                status_code=401,
                headers={
                    "WWW-Authenticate": f'Bearer resource_metadata="{resource_metadata_url}"',
                },
            )
            await response(scope, receive, send)
            return

        token = auth_header[7:]

        # --- JWT → API key resolution ---
        if token.count(".") == 2:
            if not _INTROSPECT_CLIENT_ID:
                response = JSONResponse(
                    {"error": "MCP OAuth not configured — set MCP_INTROSPECTION_CLIENT_ID"},
                    status_code=501,
                )
                await response(scope, receive, send)
                return

            api_key = await _resolve_jwt_to_api_key(token)
            if not api_key:
                response = JSONResponse(
                    {"error": "Invalid or expired MCP OAuth token"},
                    status_code=401,
                )
                await response(scope, receive, send)
                return
            set_api_key(api_key)
        else:
            set_api_key(token)

        # --- Start session manager (once) and forward to MCP handler ---
        await self._ensure_started()
        await self._handler(scope, receive, send)


def create_remote_app():
    """Create MCP ASGI app with bearer auth for remote HTTP access.

    Accepts both API keys (direct) and PropelAuth MCP OAuth 2.1 tokens.

    Builds the MCP handler directly (bypassing Starlette's lifespan wrapper)
    because FastAPI's ``app.mount()`` does not forward lifespan events to
    mounted sub-applications. The session manager is started lazily on the
    first HTTP request via a background task.
    """
    server = FastMCP("graph8")
    client = G8Client(require_key=False)
    _register_tools(server, client, mode=MCP_MODE)

    session_manager = StreamableHTTPSessionManager(
        app=server._mcp_server,
        event_store=server._event_store,
        json_response=server.settings.json_response,
        stateless=True,
    )
    handler = StreamableHTTPASGIApp(session_manager)

    return _RemoteMCPApp(handler, session_manager)


def main() -> None:
    """Entry point for `uvx g8-mcp-server` (local stdio mode)."""
    try:
        server = create_server()
    except G8ClientError as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)
    server.run(transport="stdio")
