"""FastAPI middleware for JWT validation and usage tracking."""

import time
from typing import TYPE_CHECKING

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse, Response

from .exceptions import AuthenticationError

if TYPE_CHECKING:
    from .client import VikdataClient
    from .usage import UsageReporter


class JWTValidationMiddleware(BaseHTTPMiddleware):
    """Middleware for JWT token validation."""

    def __init__(
        self,
        app: any,
        client: "VikdataClient",
        exclude_paths: list[str] | None = None,
    ) -> None:
        """Initialize JWT validation middleware.

        Args:
            app: FastAPI application
            client: VikdataClient instance
            exclude_paths: Paths to exclude from validation (default: ["/health", "/metrics"])
        """
        super().__init__(app)
        self.client = client
        self.exclude_paths = exclude_paths or ["/health", "/metrics", "/docs", "/openapi.json"]

    async def dispatch(self, request: Request, call_next: any) -> Response:
        """Process request with JWT validation.

        Args:
            request: Incoming request
            call_next: Next middleware/handler

        Returns:
            Response
        """
        # Skip excluded paths
        if any(request.url.path.startswith(p) for p in self.exclude_paths):
            return await call_next(request)

        # Extract token from Authorization header
        auth_header = request.headers.get("Authorization")
        if not auth_header or not auth_header.startswith("Bearer "):
            return JSONResponse(
                {"error": "Missing or invalid Authorization header"},
                status_code=401,
            )

        token = auth_header[7:]  # Remove "Bearer " prefix

        try:
            # Validate token and extract claims
            claims = await self.client.validate_token(token)

            # Store claims in request state for downstream use
            request.state.user = claims
            request.state.tenant_id = claims.tenant_id
            request.state.user_id = claims.sub
            request.state.permissions = claims.permissions
            if claims.roles:
                request.state.roles = claims.roles

        except AuthenticationError as e:
            return JSONResponse(
                {"error": str(e)},
                status_code=401,
            )

        return await call_next(request)


class VikdataMiddleware(BaseHTTPMiddleware):
    """Combined middleware for JWT validation and usage tracking."""

    def __init__(
        self,
        app: any,
        client: "VikdataClient",
        usage_reporter: "UsageReporter | None" = None,
        exclude_paths: list[str] | None = None,
    ) -> None:
        """Initialize vikdata middleware.

        Args:
            app: FastAPI application
            client: VikdataClient instance
            usage_reporter: Optional UsageReporter for tracking metrics
            exclude_paths: Paths to exclude from JWT validation
        """
        super().__init__(app)
        self.client = client
        self.usage_reporter = usage_reporter
        self.exclude_paths = exclude_paths or ["/health", "/metrics", "/docs", "/openapi.json"]

    async def dispatch(self, request: Request, call_next: any) -> Response:
        """Process request with JWT validation and usage tracking.

        Args:
            request: Incoming request
            call_next: Next middleware/handler

        Returns:
            Response
        """
        start_time = time.time()

        # Skip excluded paths for JWT validation
        skip_auth = any(request.url.path.startswith(p) for p in self.exclude_paths)

        if not skip_auth:
            # Extract and validate token
            auth_header = request.headers.get("Authorization")
            if not auth_header or not auth_header.startswith("Bearer "):
                return JSONResponse(
                    {"error": "Missing or invalid Authorization header"},
                    status_code=401,
                )

            token = auth_header[7:]

            try:
                claims = await self.client.validate_token(token)
                request.state.user = claims
                request.state.tenant_id = claims.tenant_id
                request.state.user_id = claims.sub
                request.state.permissions = claims.permissions
                if claims.roles:
                    request.state.roles = claims.roles

            except AuthenticationError as e:
                return JSONResponse(
                    {"error": str(e)},
                    status_code=401,
                )

        # Process request
        response = await call_next(request)

        # Track usage if reporter is configured and we have tenant_id
        if self.usage_reporter and hasattr(request.state, "tenant_id"):
            latency_ms = (time.time() - start_time) * 1000
            await self.usage_reporter.record(
                tenant_id=request.state.tenant_id,
                endpoint=request.url.path,
                method=request.method,
                latency_ms=latency_ms,
                status_code=response.status_code,
            )

        return response


class UsageMiddleware(BaseHTTPMiddleware):
    """Middleware for usage tracking only (no JWT validation)."""

    def __init__(
        self,
        app: any,
        usage_reporter: "UsageReporter",
    ) -> None:
        """Initialize usage tracking middleware.

        Args:
            app: FastAPI application
            usage_reporter: UsageReporter instance
        """
        super().__init__(app)
        self.usage_reporter = usage_reporter

    async def dispatch(self, request: Request, call_next: any) -> Response:
        """Process request with usage tracking.

        Args:
            request: Incoming request
            call_next: Next middleware/handler

        Returns:
            Response
        """
        start_time = time.time()

        # Process request
        response = await call_next(request)

        # Track usage if tenant_id is available
        if hasattr(request.state, "tenant_id"):
            latency_ms = (time.time() - start_time) * 1000
            await self.usage_reporter.record(
                tenant_id=request.state.tenant_id,
                endpoint=request.url.path,
                method=request.method,
                latency_ms=latency_ms,
                status_code=response.status_code,
            )

        return response
