"""Crustdata MCP Server - B2B data platform API access for AI assistants."""

import os
from contextlib import asynccontextmanager
from copy import copy
from dataclasses import dataclass
from typing import AsyncIterator

import uvicorn
from mcp.server.auth.settings import AuthSettings, ClientRegistrationOptions, RevocationOptions
from mcp.server.fastmcp import FastMCP
from mcp.server.transport_security import TransportSecuritySettings
from pydantic import AnyHttpUrl
from starlette.routing import Route

from .client import CrustdataClient
from .client_factory import ClientFactory, get_client_factory, get_client_for_request, set_request_token
from .config import config, get_token
from .oauth import CrustdataOAuthProvider


@dataclass
class AppContext:
    """Application context available during server lifespan.

    Contains the ClientFactory which manages CrustdataClient instances.
    For HTTP transport, each request gets a client based on the provided token.
    For STDIO transport, a single client is used based on the env var token.
    """

    client_factory: ClientFactory

    @property
    def client(self) -> CrustdataClient:
        """Get client for current request (backward compatibility).

        Deprecated: Use get_client_for_request() directly in tools instead.
        """
        return get_client_for_request()


@asynccontextmanager
async def app_lifespan_stdio(server: FastMCP) -> AsyncIterator[AppContext]:
    """Manage application lifecycle for STDIO transport.

    Validates that CRUSTDATA_TOKEN is set at startup, since STDIO mode
    requires a pre-configured token (single-tenant mode).

    Yields:
        AppContext with the client factory.
    """
    api_token = get_token()
    if not api_token:
        raise ValueError(
            "CRUSTDATA_TOKEN environment variable is required. "
            "Set it to your Crustdata API token."
        )

    factory = get_client_factory()
    factory.set_transport_mode("stdio")
    # Pre-validate by creating the STDIO client
    _ = factory._get_stdio_client()

    try:
        yield AppContext(client_factory=factory)
    finally:
        await factory.close_all()


@asynccontextmanager
async def app_lifespan_http(server: FastMCP) -> AsyncIterator[AppContext]:
    """Manage application lifecycle for HTTP transport.

    Does NOT require CRUSTDATA_TOKEN at startup. Each HTTP request
    provides its own token via the Authorization header (multi-tenant mode).

    Note: Setting transport mode to "http" ensures that if a request reaches
    tool code without an AccessToken, it will fail explicitly rather than
    silently falling back to the CRUSTDATA_TOKEN env var.

    Yields:
        AppContext with the client factory.
    """
    factory = get_client_factory()
    factory.set_transport_mode("http")
    try:
        yield AppContext(client_factory=factory)
    finally:
        await factory.close_all()


# Backward compatibility alias
app_lifespan = app_lifespan_stdio


# Create the MCP server instance
mcp = FastMCP(
    name="Crustdata MCP Server",
    instructions="""
    Crustdata MCP Server provides access to B2B company and people data.

    Use this server to:
    - Search and identify companies by name, domain, or LinkedIn URL
    - Enrich company profiles with detailed business data
    - Search for people/decision makers at companies
    - Monitor companies with watchers for job postings, funding, etc.
    - Search and fetch web content

    All tools require a valid Crustdata API token set in the CRUSTDATA_TOKEN
    environment variable.

    Credit costs vary by endpoint - check your credit balance with the
    crustdata_credits_check tool.
    """,
    lifespan=app_lifespan,
    log_level=config.log_level,  # type: ignore[arg-type]
)


# Import tools to register them with the server
# These imports must come after mcp is defined to avoid circular imports
from . import tools  # noqa: E402, F401


def main() -> None:
    """Run the MCP server with stdio transport."""
    mcp.run(transport="stdio")


def main_http() -> None:
    """Run the MCP server with HTTP transport and OAuth authentication.

    Users authenticate via OAuth 2.1 flow — Claude.ai redirects them to a
    login page where they paste their Crustdata API key. The key is validated,
    then used as the Bearer token for all subsequent MCP requests.

    Configure via environment variables:
        CRUSTDATA_HTTP_HOST: Host to bind to (default: 127.0.0.1)
        CRUSTDATA_HTTP_PORT: Port to listen on (default: 8000)
        CRUSTDATA_HTTP_PATH: URL path for MCP endpoint (default: /mcp)
        CRUSTDATA_BASE_URL: Public base URL (default: http://host:port)
    """
    base_url = os.environ.get(
        "CRUSTDATA_BASE_URL",
        f"http://{config.http_host}:{config.http_port}",
    )
    issuer_url = AnyHttpUrl(base_url)
    resource_url = AnyHttpUrl(f"{base_url}{config.http_path}")

    oauth_provider = CrustdataOAuthProvider(base_url=base_url)

    # Disable DNS rebinding protection to allow connections from any host
    transport_security = TransportSecuritySettings(
        enable_dns_rebinding_protection=False,
    )

    mcp_http = FastMCP(
        name="Crustdata MCP Server",
        instructions=mcp.instructions,
        lifespan=app_lifespan_http,
        log_level=config.log_level,  # type: ignore[arg-type]
        stateless_http=True,
        transport_security=transport_security,
        auth=AuthSettings(
            issuer_url=issuer_url,
            resource_server_url=resource_url,
            client_registration_options=ClientRegistrationOptions(
                enabled=True,
                valid_scopes=["read", "write"],
                default_scopes=["read", "write"],
            ),
            revocation_options=RevocationOptions(enabled=True),
        ),
        auth_server_provider=oauth_provider,
    )

    # Re-register all tools from the main server (shallow copy to avoid shared state).
    # Uses FastMCP internals (mcp SDK ~1.x) — no public API exists. Revisit on SDK upgrade.
    for tool_name, tool_def in mcp._tool_manager._tools.items():
        mcp_http._tool_manager._tools[tool_name] = copy(tool_def)

    # Add custom login routes for the OAuth flow
    mcp_http._custom_starlette_routes.extend([
        Route("/oauth/login", endpoint=oauth_provider.handle_login_get, methods=["GET"]),
        Route("/oauth/login", endpoint=oauth_provider.handle_login_post, methods=["POST"]),
        Route("/oauth/auth0-callback", endpoint=oauth_provider.handle_auth0_callback, methods=["GET"]),
    ])

    # Get the ASGI app and wrap with bearer token extraction middleware
    inner_app = mcp_http.streamable_http_app()

    async def bearer_token_middleware(scope, receive, send):
        """ASGI middleware to extract Bearer token for the client factory."""
        if scope["type"] == "http":
            set_request_token(None)  # Clear stale token from keep-alive connections
            headers = dict(scope.get("headers", []))
            auth_header = headers.get(b"authorization", b"").decode()
            if auth_header.startswith("Bearer "):
                token = auth_header[7:]
                set_request_token(token)

        await inner_app(scope, receive, send)

    uvicorn.run(
        bearer_token_middleware,
        host=config.http_host,
        port=config.http_port,
        log_level=config.log_level.lower(),
        access_log=False,  # Prevent API keys in Bearer headers from leaking into logs
    )


if __name__ == "__main__":
    main()
