"""Unified authentication middleware for the Graph API Gateway.

Supports API key, JWT, and MCP token auth methods across all protocols.
"""

from __future__ import annotations

import os
import logging
from typing import Optional

from fastapi import Request, WebSocket
from fastapi.exceptions import HTTPException

from ..context import RequestContext

logger = logging.getLogger(__name__)

# Environment-based API key for initial deployment
# In production, this would validate against a key store
_GATEWAY_API_KEY = os.environ.get("DATABRIDGE_GATEWAY_API_KEY", "")


def _extract_api_key(request: Request) -> Optional[str]:
    """Extract API key from Authorization header or query param."""
    auth = request.headers.get("authorization", "")
    if auth.startswith("Bearer "):
        return auth[7:]
    return request.query_params.get("api_key")


def _extract_tenant_id(request: Request) -> str:
    """Extract tenant ID from header or query param."""
    return (
        request.headers.get("x-tenant-id")
        or request.query_params.get("tenant_id")
        or "default"
    )


def _resolve_tier(tenant_registry, tenant_id: str) -> str:
    """Resolve tenant tier from registry, defaulting to CE."""
    if tenant_registry is None:
        return "CE"
    cfg = tenant_registry.get(tenant_id)
    return cfg.tier if cfg else "CE"


async def get_request_context(
    request: Request,
    tenant_registry=None,
    protocol: str = "rest",
) -> RequestContext:
    """Build a RequestContext from an incoming HTTP request.

    Auth is optional — if no key is provided, the request proceeds
    as the default tenant with CE tier. When DATABRIDGE_GATEWAY_API_KEY
    is set, requests with an invalid key are rejected.
    """
    api_key = _extract_api_key(request)
    tenant_id = _extract_tenant_id(request)
    auth_method = "none"

    if api_key:
        if _GATEWAY_API_KEY and api_key != _GATEWAY_API_KEY:
            raise HTTPException(status_code=401, detail="Invalid API key")
        auth_method = "api_key"

    tier = _resolve_tier(tenant_registry, tenant_id)

    return RequestContext(
        tenant_id=tenant_id,
        tenant_tier=tier,
        auth_method=auth_method,
        protocol=protocol,
    )


async def authenticate_ws(websocket: WebSocket, tenant_registry=None) -> RequestContext:
    """Authenticate a WebSocket connection.

    Expects the first message to be a JSON auth payload:
    {"token": "...", "tenant_id": "..."}

    If no auth message is received, proceeds as default tenant.
    """
    try:
        auth_data = await websocket.receive_json()
    except Exception:
        return RequestContext(protocol="ws")

    token = auth_data.get("token", "")
    tenant_id = auth_data.get("tenant_id", "default")
    auth_method = "none"

    if token:
        if _GATEWAY_API_KEY and token != _GATEWAY_API_KEY:
            await websocket.close(code=4001, reason="Invalid token")
            raise HTTPException(status_code=401, detail="Invalid WebSocket token")
        auth_method = "api_key"

    tier = _resolve_tier(tenant_registry, tenant_id)

    return RequestContext(
        tenant_id=tenant_id,
        tenant_tier=tier,
        auth_method=auth_method,
        protocol="ws",
    )
