"""
Authentication providers for fabias.

Supports multiple authentication methods:
- ServicePrincipalAuth: OAuth2 client credentials flow for standalone usage
- FabricAuth: Uses notebookutils when running inside Fabric notebooks
- AutoAuth: Automatically selects the best auth method based on environment
"""

from abc import ABC, abstractmethod
from datetime import datetime, timedelta, timezone
from typing import Callable, Dict, Optional, Tuple, cast

import requests

from .exceptions import AuthenticationError
from .runtime import runtime


class AuthProvider(ABC):
    """
    Abstract base class for authentication providers.

    All authentication providers must implement the `getToken` method
    to obtain an access token for a given scope/resource.
    """

    @abstractmethod
    def getToken(self, scope: str) -> str:
        """
        Get an access token for the specified scope.

        Args:
            scope: OAuth2 scope (e.g., "https://api.fabric.microsoft.com/.default")

        Returns:
            str: Bearer access token

        Raises:
            AuthenticationError: If token acquisition fails
        """
        pass

    def getAuthorizationHeader(self, scope: str) -> str:
        """
        Get the Authorization header value for API requests.

        Args:
            scope: OAuth2 scope for the target API

        Returns:
            str: Authorization header value (e.g., "Bearer <token>")
        """
        return f"Bearer {self.getToken(scope)}"


class ServicePrincipalAuth(AuthProvider):
    """
    OAuth2 client credentials authentication using service principal.

    This provider authenticates using Azure AD client credentials flow,
    suitable for service-to-service authentication without user interaction.

    Tokens are cached and automatically refreshed when expired.

    Attributes:
        tenant_id (str): Azure AD tenant ID
        client_id (str): Application (client) ID
        client_secret (str): Client secret value

    Examples:
        >>> auth = ServicePrincipalAuth(
        ...     tenant_id="your-tenant-id",
        ...     client_id="your-client-id",
        ...     client_secret="your-secret"
        ... )
        >>> token = auth.getToken("https://api.fabric.microsoft.com/.default")
    """

    def __init__(self, tenant_id: str, client_id: str, client_secret: str):
        """
        Initialize the service principal authentication provider.

        Args:
            tenant_id: Azure AD tenant ID (GUID or domain)
            client_id: Application (client) ID from Azure AD
            client_secret: Client secret value
        """
        self._tenant_id = tenant_id
        self._client_id = client_id
        self._client_secret = client_secret

        # Token cache: scope -> (token, expiry)
        self._token_cache: Dict[str, tuple] = {}

    def getToken(self, scope: str) -> str:
        """
        Get an access token for the specified scope.

        Retrieves a cached token if still valid, otherwise acquires a new
        token using the OAuth2 client credentials flow.

        Args:
            scope: OAuth2 scope (e.g., "https://api.fabric.microsoft.com/.default")

        Returns:
            str: Access token (without "Bearer " prefix)

        Raises:
            AuthenticationError: If token acquisition fails
        """
        # Check cache
        if scope in self._token_cache:
            token, expiry = self._token_cache[scope]
            # Add buffer of 60 seconds before expiry
            if expiry > datetime.now(timezone.utc) + timedelta(seconds=60):
                return token

        # Acquire new token
        token, expiry = self._acquire_token(scope)
        self._token_cache[scope] = (token, expiry)
        return token

    def _acquire_token(self, scope: str) -> Tuple[str, datetime]:
        """
        Acquire a new token from Azure AD.

        Args:
            scope: OAuth2 scope

        Returns:
            Tuple[str, datetime]: (access_token, expiry_datetime)

        Raises:
            AuthenticationError: If token acquisition fails
        """
        token_url = f"https://login.microsoftonline.com/{self._tenant_id}/oauth2/v2.0/token"

        headers = {"Content-Type": "application/x-www-form-urlencoded"}

        data = {
            "grant_type": "client_credentials",
            "scope": scope,
            "client_id": self._client_id,
            "client_secret": self._client_secret,
        }

        try:
            response = requests.post(token_url, data=data, headers=headers)

            if response.status_code != 200:
                error_data = (
                    response.json()
                    if response.headers.get("content-type", "").startswith("application/json")
                    else {}
                )
                error_msg = error_data.get(
                    "error_description", error_data.get("error", response.text)
                )
                raise AuthenticationError(f"Token acquisition failed: {error_msg}")

            token_data = response.json()
            access_token = token_data.get("access_token")
            expires_in: int = token_data.get("expires_in", 3600)

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

            expiry = datetime.now(timezone.utc) + timedelta(seconds=expires_in)

            return str(access_token), expiry

        except requests.RequestException as e:
            raise AuthenticationError("Network error during authentication", cause=e)


class RefreshTokenAuth(AuthProvider):
    """
    User-delegated authentication using refresh tokens.

    This provider supports two modes:

    **Self-managed mode**: You provide a refresh_token, and the provider
    handles token rotation automatically. Use `on_refresh` callback to
    persist the new refresh token (e.g., to Key Vault).

    **Delegated mode**: You provide a `token_provider` callable that
    handles all token acquisition externally (e.g., via an Azure Function
    token broker). The provider just calls your function.

    Tokens are cached and automatically refreshed when expired.

    Attributes:
        tenant_id (str): Azure AD tenant ID
        client_id (str): Application (client) ID

    Examples:
        Self-managed with Key Vault persistence:

        >>> from fabias.integrations import keyvault
        >>> auth = RefreshTokenAuth(
        ...     tenant_id="your-tenant-id",
        ...     client_id="your-client-id",
        ...     refresh_token=keyvault.get("graph-refresh-token"),
        ...     on_refresh=lambda token: keyvault.set("graph-refresh-token", token)
        ... )
        >>> token = auth.getToken("https://graph.microsoft.com/.default")

        Delegated mode with external token broker:

        >>> def get_from_broker(scope: str) -> tuple[str, datetime]:
        ...     response = requests.post("https://my-function/api/token", json={"scope": scope})
        ...     data = response.json()
        ...     return data["access_token"], datetime.fromisoformat(data["expires_at"])
        ...
        >>> auth = RefreshTokenAuth(
        ...     tenant_id="your-tenant-id",
        ...     client_id="your-client-id",
        ...     token_provider=get_from_broker
        ... )

    Warning:
        Self-managed mode is NOT safe for concurrent use across multiple jobs
        that might refresh the token simultaneously. For production multi-job
        scenarios, use delegated mode with a token broker service.
    """

    def __init__(
        self,
        tenant_id: str,
        client_id: str,
        refresh_token: Optional[str] = None,
        on_refresh: Optional["Callable[[str], None]"] = None,
        token_provider: Optional["Callable[[str], Tuple[str, datetime]]"] = None,
    ):
        """
        Initialize the refresh token authentication provider.

        Args:
            tenant_id: Azure AD tenant ID (GUID or domain)
            client_id: Application (client) ID from Azure AD
            refresh_token: Initial refresh token (for self-managed mode)
            on_refresh: Callback to persist new refresh token (for self-managed mode).
                        Called with the new refresh_token string after each refresh.
            token_provider: External function that returns (access_token, expiry)
                           for a given scope (for delegated mode). If provided,
                           refresh_token and on_refresh are ignored.

        Raises:
            ValueError: If neither refresh_token nor token_provider is provided
        """
        if not refresh_token and not token_provider:
            raise ValueError(
                "Either refresh_token (self-managed) or token_provider (delegated) is required"
            )

        self._tenant_id = tenant_id
        self._client_id = client_id
        self._refresh_token = refresh_token
        self._on_refresh = on_refresh
        self._token_provider = token_provider

        # Token cache: scope -> (token, expiry)
        self._token_cache: Dict[str, tuple] = {}

    @property
    def refresh_token(self) -> Optional[str]:
        """Get the current refresh token (may have been rotated)."""
        return self._refresh_token

    def getToken(self, scope: str) -> str:
        """
        Get an access token for the specified scope.

        Retrieves a cached token if still valid, otherwise acquires a new
        token using either the token_provider or refresh token flow.

        Args:
            scope: OAuth2 scope (e.g., "https://graph.microsoft.com/.default")

        Returns:
            str: Access token (without "Bearer " prefix)

        Raises:
            AuthenticationError: If token acquisition fails
        """
        # Check cache
        if scope in self._token_cache:
            token, expiry = self._token_cache[scope]
            # Add buffer of 60 seconds before expiry
            if expiry > datetime.now(timezone.utc) + timedelta(seconds=60):
                return token

        # Acquire new token
        token, expiry = self._acquire_token(scope)
        self._token_cache[scope] = (token, expiry)
        return token

    def _acquire_token(self, scope: str) -> Tuple[str, datetime]:
        """
        Acquire a new access token.

        Uses token_provider if available, otherwise uses refresh token flow.

        Args:
            scope: OAuth2 scope

        Returns:
            Tuple[str, datetime]: (access_token, expiry_datetime)

        Raises:
            AuthenticationError: If token acquisition fails
        """
        # Delegated mode - use external provider
        if self._token_provider:
            try:
                return self._token_provider(scope)
            except Exception as e:
                raise AuthenticationError(f"Token provider failed: {str(e)}", cause=e)

        # Self-managed mode - use refresh token
        return self._do_refresh(scope)

    def _do_refresh(self, scope: str) -> Tuple[str, datetime]:
        """
        Exchange refresh token for new access token.

        Also handles token rotation - the new refresh token is passed
        to the on_refresh callback if provided.

        Args:
            scope: OAuth2 scope

        Returns:
            Tuple[str, datetime]: (access_token, expiry_datetime)

        Raises:
            AuthenticationError: If refresh fails
        """
        if not self._refresh_token:
            raise AuthenticationError("No refresh token available")

        token_url = f"https://login.microsoftonline.com/{self._tenant_id}/oauth2/v2.0/token"

        headers = {"Content-Type": "application/x-www-form-urlencoded"}

        data = {
            "grant_type": "refresh_token",
            "scope": scope,
            "client_id": self._client_id,
            "refresh_token": self._refresh_token,
        }

        try:
            response = requests.post(token_url, data=data, headers=headers)

            if response.status_code != 200:
                error_data = (
                    response.json()
                    if response.headers.get("content-type", "").startswith("application/json")
                    else {}
                )
                error = error_data.get("error", "unknown")
                error_msg = error_data.get(
                    "error_description", error_data.get("error", response.text)
                )

                # Check for invalid_grant - token may have been rotated by another process
                if error == "invalid_grant":
                    raise AuthenticationError(
                        f"Refresh token invalid or expired. {error_msg}. "
                        "The token may have been rotated by another process, "
                        "or it has expired (90 day limit)."
                    )

                raise AuthenticationError(f"Token refresh failed: {error_msg}")

            token_data = response.json()
            access_token = token_data.get("access_token")
            expires_in: int = token_data.get("expires_in", 3600)
            new_refresh_token = token_data.get("refresh_token")

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

            # Update stored refresh token (it rotates with each use)
            if new_refresh_token:
                self._refresh_token = new_refresh_token

                # Call persistence callback if provided
                if self._on_refresh:
                    try:
                        self._on_refresh(new_refresh_token)
                    except Exception as e:
                        # Log but don't fail - the token worked, persistence is secondary
                        import logging

                        logging.warning(f"Failed to persist refresh token: {e}")

            expiry = datetime.now(timezone.utc) + timedelta(seconds=expires_in)

            return str(access_token), expiry

        except requests.RequestException as e:
            raise AuthenticationError("Network error during token refresh", cause=e)


class FabricAuth(AuthProvider):
    """
    Authentication provider using Fabric's notebookutils.

    This provider leverages the notebookutils.credentials module available
    in Microsoft Fabric notebooks to obtain tokens without requiring
    explicit credentials.

    Only works when running inside a Fabric notebook environment.

    Examples:
        >>> from fabias import runtime
        >>> if runtime.isFabric:
        ...     auth = FabricAuth()
        ...     token = auth.getToken("https://api.fabric.microsoft.com/.default")
    """

    def __init__(self):
        """
        Initialize the Fabric authentication provider.

        Raises:
            AuthenticationError: If not running inside Fabric environment
        """
        if not runtime.isFabric:
            raise AuthenticationError(
                "FabricAuth requires Fabric notebook environment. "
                "Use ServicePrincipalAuth for standalone execution."
            )

        self._notebookutils = runtime.notebookutils

    def getToken(self, scope: str) -> str:
        """
        Get an access token using notebookutils.

        Args:
            scope: OAuth2 scope (e.g., "https://api.fabric.microsoft.com/.default")

        Returns:
            str: Access token

        Raises:
            AuthenticationError: If token acquisition fails
        """
        try:
            # notebookutils.credentials.getToken expects a resource URL
            # Convert scope to resource if needed (remove /.default suffix)
            resource = scope.replace("/.default", "")

            token = self._notebookutils.credentials.getToken(resource)

            if not token:
                raise AuthenticationError(f"No token returned for scope: {scope}")

            return cast(str, token)

        except AttributeError as e:
            raise AuthenticationError("notebookutils.credentials.getToken not available", cause=e)
        except Exception as e:
            raise AuthenticationError(f"Failed to get token from notebookutils: {str(e)}", cause=e)


class AutoAuth(AuthProvider):
    """
    Automatic authentication provider that selects the best method.

    This provider automatically chooses between FabricAuth (when inside
    Fabric) and ServicePrincipalAuth (when standalone) based on the
    detected runtime environment.

    For standalone usage, credentials must be provided either directly
    or via environment variables/configuration.

    Examples:
        Inside Fabric:
        >>> auth = AutoAuth()  # Automatically uses FabricAuth

        Standalone with explicit credentials:
        >>> auth = AutoAuth(
        ...     tenant_id="...",
        ...     client_id="...",
        ...     client_secret="..."
        ... )
    """

    def __init__(
        self,
        tenant_id: Optional[str] = None,
        client_id: Optional[str] = None,
        client_secret: Optional[str] = None,
    ):
        """
        Initialize with optional service principal credentials.

        If running in Fabric, credentials are ignored. If running standalone,
        credentials are required.

        Args:
            tenant_id: Azure AD tenant ID (for standalone)
            client_id: Application client ID (for standalone)
            client_secret: Client secret (for standalone)
        """
        self._provider: Optional[AuthProvider] = None

        if runtime.isFabric:
            self._provider = FabricAuth()
        else:
            if not all([tenant_id, client_id, client_secret]):
                raise AuthenticationError(
                    "Service principal credentials required for standalone execution. "
                    "Provide tenant_id, client_id, and client_secret."
                )
            # Type narrowing: we've already validated these are not None
            self._provider = ServicePrincipalAuth(
                tenant_id=cast(str, tenant_id),
                client_id=cast(str, client_id),
                client_secret=cast(str, client_secret),
            )

    def getToken(self, scope: str) -> str:
        """
        Get an access token using the automatically selected provider.

        Args:
            scope: OAuth2 scope

        Returns:
            str: Access token
        """
        assert self._provider is not None, "Provider not initialized"
        return self._provider.getToken(scope)
