"""Authentication middleware with multi-tenant support."""

from collections.abc import Callable

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

from pydanticai_multiagent.services.tenant import Tenant, TenantService


class AuthMiddleware(BaseHTTPMiddleware):
    """Middleware for authenticating requests.

    Supports two authentication modes:
    1. Tenant API keys (sk-tenant-xxx) - For platform clients
    2. JWT tokens - For end users (optional)

    Example usage in app.py:
        ```python
        app.add_middleware(
            AuthMiddleware,
            secret_key="your-secret",
            require_auth=True,  # Reject unauthenticated requests
        )
        ```
    """

    def __init__(
        self,
        app: Callable,
        secret_key: str,
        exclude_paths: list[str] | None = None,
        require_auth: bool = False,
    ) -> None:
        """Initialize auth middleware.

        Args:
            app: The FastAPI application.
            secret_key: Secret key for JWT validation.
            exclude_paths: Paths to exclude from auth.
            require_auth: If True, reject requests without valid auth.
        """
        super().__init__(app)
        self.secret_key = secret_key
        self.exclude_paths = exclude_paths or [
            "/health",
            "/ready",
            "/docs",
            "/openapi.json",
            "/redoc",
        ]
        self.require_auth = require_auth

    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        """Process the request through authentication."""
        # Skip auth for excluded paths
        if any(request.url.path.startswith(path) for path in self.exclude_paths):
            return await call_next(request)

        # Extract token from header
        auth_header = request.headers.get("Authorization")

        if not auth_header:
            if self.require_auth:
                return JSONResponse(
                    status_code=401,
                    content={"error": "Authorization header required"},
                )
            request.state.tenant = None
            request.state.user_id = None
            request.state.user_roles = []
            return await call_next(request)

        # Parse Bearer token
        try:
            scheme, token = auth_header.split()
            if scheme.lower() != "bearer":
                if self.require_auth:
                    return JSONResponse(
                        status_code=401,
                        content={"error": "Invalid authorization scheme"},
                    )
                request.state.tenant = None
                request.state.user_id = None
                request.state.user_roles = []
                return await call_next(request)
        except ValueError:
            if self.require_auth:
                return JSONResponse(
                    status_code=401,
                    content={"error": "Invalid authorization header"},
                )
            request.state.tenant = None
            request.state.user_id = None
            request.state.user_roles = []
            return await call_next(request)

        # Check if this is a tenant API key
        if token.startswith("sk-tenant-"):
            tenant = await self._validate_tenant_key(request, token)
            if tenant:
                request.state.tenant = tenant
                request.state.user_id = f"tenant:{tenant.id}"
                request.state.user_roles = ["tenant"]
            elif self.require_auth:
                return JSONResponse(
                    status_code=401,
                    content={"error": "Invalid API key"},
                )
            else:
                request.state.tenant = None
                request.state.user_id = None
                request.state.user_roles = []
        else:
            # Try JWT validation
            user_info = self._validate_jwt(token)
            request.state.tenant = None
            if user_info:
                request.state.user_id = user_info.get("user_id")
                request.state.user_roles = user_info.get("roles", [])
            elif self.require_auth:
                return JSONResponse(
                    status_code=401,
                    content={"error": "Invalid token"},
                )
            else:
                request.state.user_id = None
                request.state.user_roles = []

        return await call_next(request)

    async def _validate_tenant_key(
        self, request: Request, api_key: str
    ) -> Tenant | None:
        """Validate a tenant API key.

        Args:
            request: The incoming request (to access app.state.db).
            api_key: The tenant API key.

        Returns:
            Tenant object or None if invalid.
        """
        db = getattr(request.app.state, "db", None)
        if not db:
            return None

        tenant_service = TenantService(db)
        return await tenant_service.get_tenant_by_api_key(api_key)

    def _validate_jwt(self, token: str) -> dict | None:
        """Validate a JWT token and return user info.

        Args:
            token: The JWT token.

        Returns:
            User info dict or None if invalid.
        """
        # Placeholder - implement actual JWT validation
        # Example with PyJWT:
        # try:
        #     payload = jwt.decode(token, self.secret_key, algorithms=["HS256"])
        #     return {
        #         "user_id": payload["sub"],
        #         "roles": payload.get("roles", []),
        #     }
        # except jwt.InvalidTokenError:
        #     return None

        return None
