"""
Authentication module for Credo AI SDK.

This module provides authentication handlers that support both API key and Bearer token
authentication, with automatic JWT token exchange for API key users.
"""

from typing import Dict, Literal, Optional

import httpx


class AuthenticationError(Exception):
    """Raised when authentication fails"""

    def __init__(
        self,
        message: str,
        status_code: Optional[int] = None,
        cause: Optional[Exception] = None,
    ):
        super().__init__(message)
        self.status_code = status_code
        self.cause = cause


class CredoAIAuth:
    """
    Authentication handler for CredoAI SDK.

    Supports both API key and Bearer token authentication:
    - API key authentication: Exchanges API key for JWT token via server auth endpoint
    - Bearer token authentication: Uses provided JWT token directly

    All authenticated requests use Bearer token format internally for consistency.
    """

    def __init__(self, base_url: str, tenant: str = "default", timeout: float = 30.0):
        """
        Initialize authentication handler.

        Args:
            base_url: Base URL of the API service
            tenant: Tenant identifier for authentication context
            timeout: Timeout for authentication requests
        """
        self.base_url = base_url.rstrip("/")
        self.tenant = tenant
        self.timeout = timeout
        self._access_token: Optional[str] = None
        self._auth_type: Optional[Literal["api_key", "bearer"]] = None

    def authenticate_with_api_key(self, api_key: str) -> str:
        """
        Exchange API key for JWT token via server's authentication endpoint.

        This method makes a synchronous request to the server's auth endpoint to
        exchange the provided API key for a JWT access token.

        Args:
            api_key: API key to exchange for JWT token

        Returns:
            str: JWT access token

        Raises:
            AuthenticationError: If API key exchange fails
        """
        if not api_key:
            raise AuthenticationError("API key cannot be empty")

        # Prepare authentication request
        auth_url = f"{self.base_url}/auth/token"
        headers = {
            "X-API-Key": api_key,
            "X-Tenant": self.tenant,
            "Accept": "application/json",
            "Content-Type": "application/json",
        }

        try:
            with httpx.Client(timeout=self.timeout) as client:
                response = client.post(auth_url, headers=headers)

                if response.status_code == 401:
                    raise AuthenticationError(
                        f"Invalid API key for tenant '{self.tenant}'", status_code=401
                    )
                elif response.status_code == 403:
                    raise AuthenticationError(
                        f"API key access denied for tenant '{self.tenant}'",
                        status_code=403,
                    )
                elif response.status_code >= 400:
                    raise AuthenticationError(
                        f"Authentication failed: {response.status_code} {response.text}",
                        status_code=response.status_code,
                    )

                # Parse response to extract access token
                try:
                    auth_data = response.json()
                    access_token = auth_data.get("access_token")

                    if not access_token:
                        raise AuthenticationError(
                            "No access token in authentication response"
                        )

                    self._access_token = access_token
                    self._auth_type = "api_key"
                    return access_token

                except (ValueError, KeyError) as e:
                    raise AuthenticationError(
                        f"Invalid authentication response format: {e}", cause=e
                    )

        except httpx.RequestError as e:
            raise AuthenticationError(
                f"Authentication service unavailable: {e}", cause=e
            )

    async def authenticate_with_api_key_async(self, api_key: str) -> str:
        """
        Async version of API key authentication.

        Args:
            api_key: API key to exchange for JWT token

        Returns:
            str: JWT access token

        Raises:
            AuthenticationError: If API key exchange fails
        """
        if not api_key:
            raise AuthenticationError("API key cannot be empty")

        # Prepare authentication request
        auth_url = f"{self.base_url}/auth/token"
        headers = {
            "X-API-Key": api_key,
            "X-Tenant": self.tenant,
            "Accept": "application/json",
            "Content-Type": "application/json",
        }

        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.post(auth_url, headers=headers)

                if response.status_code == 401:
                    raise AuthenticationError(
                        f"Invalid API key for tenant '{self.tenant}'", status_code=401
                    )
                elif response.status_code == 403:
                    raise AuthenticationError(
                        f"API key access denied for tenant '{self.tenant}'",
                        status_code=403,
                    )
                elif response.status_code >= 400:
                    raise AuthenticationError(
                        f"Authentication failed: {response.status_code} {response.text}",
                        status_code=response.status_code,
                    )

                # Parse response to extract access token
                try:
                    auth_data = response.json()
                    access_token = auth_data.get("access_token")

                    if not access_token:
                        raise AuthenticationError(
                            "No access token in authentication response"
                        )

                    self._access_token = access_token
                    self._auth_type = "api_key"
                    return access_token

                except (ValueError, KeyError) as e:
                    raise AuthenticationError(
                        f"Invalid authentication response format: {e}", cause=e
                    )

        except httpx.RequestError as e:
            raise AuthenticationError(
                f"Authentication service unavailable: {e}", cause=e
            )

    def authenticate_with_bearer(self, bearer_token: str) -> str:
        """
        Use provided bearer token directly for authentication.

        This method validates the token format and stores it for use in
        authenticated requests.

        Args:
            bearer_token: JWT bearer token

        Returns:
            str: The same bearer token (for consistency with API key flow)

        Raises:
            AuthenticationError: If bearer token is invalid
        """
        if not bearer_token:
            raise AuthenticationError("Bearer token cannot be empty")

        # Basic validation - ensure it looks like a JWT token
        if not bearer_token.count(".") >= 2:
            raise AuthenticationError("Bearer token does not appear to be a valid JWT")

        self._access_token = bearer_token
        self._auth_type = "bearer"
        return bearer_token

    def get_auth_headers(self) -> Dict[str, str]:
        """
        Get authentication headers for API requests.

        Returns:
            Dict[str, str]: Headers containing Authorization bearer token

        Raises:
            AuthenticationError: If no valid authentication is available
        """
        if not self._access_token:
            raise AuthenticationError(
                "No valid authentication. Call authenticate_with_api_key() or "
                "authenticate_with_bearer() first."
            )

        return {
            "Authorization": f"Bearer {self._access_token}",
            "X-Tenant": self.tenant,
        }

    def is_authenticated(self) -> bool:
        """
        Check if authentication is available.

        Returns:
            bool: True if authenticated, False otherwise
        """
        return self._access_token is not None

    def get_auth_type(self) -> Optional[Literal["api_key", "bearer"]]:
        """
        Get the type of authentication currently in use.

        Returns:
            Optional[Literal["api_key", "bearer"]]: Authentication type or None if not authenticated
        """
        return self._auth_type

    def clear_authentication(self) -> None:
        """Clear stored authentication data."""
        self._access_token = None
        self._auth_type = None


def create_auth_handler(
    base_url: str,
    auth_type: Literal["api_key", "bearer"],
    token: str,
    tenant: str = "default",
    timeout: float = 30.0,
) -> CredoAIAuth:
    """
    Create and configure an authentication handler.

    Args:
        base_url: Base URL of the API service
        auth_type: Type of authentication ("api_key" or "bearer")
        token: API key or bearer token
        tenant: Tenant identifier
        timeout: Timeout for authentication requests

    Returns:
        CredoAIAuth: Configured authentication handler

    Raises:
        AuthenticationError: If authentication setup fails
    """
    auth_handler = CredoAIAuth(base_url, tenant, timeout)

    if auth_type == "api_key":
        auth_handler.authenticate_with_api_key(token)
    elif auth_type == "bearer":
        auth_handler.authenticate_with_bearer(token)
    else:
        raise ValueError(f"Unsupported auth_type: {auth_type}")

    return auth_handler


async def create_auth_handler_async(
    base_url: str,
    auth_type: Literal["api_key", "bearer"],
    token: str,
    tenant: str = "default",
    timeout: float = 30.0,
) -> CredoAIAuth:
    """
    Create and configure an authentication handler asynchronously.

    Args:
        base_url: Base URL of the API service
        auth_type: Type of authentication ("api_key" or "bearer")
        token: API key or bearer token
        tenant: Tenant identifier
        timeout: Timeout for authentication requests

    Returns:
        CredoAIAuth: Configured authentication handler

    Raises:
        AuthenticationError: If authentication setup fails
    """
    auth_handler = CredoAIAuth(base_url, tenant, timeout)

    if auth_type == "api_key":
        await auth_handler.authenticate_with_api_key_async(token)
    elif auth_type == "bearer":
        auth_handler.authenticate_with_bearer(token)
    else:
        raise ValueError(f"Unsupported auth_type: {auth_type}")

    return auth_handler
