"""M2M Token Manager for Impersonation."""
import datetime
from time import sleep
from typing import Annotated, Any, Self

import httpx
from pydantic import (
    BaseModel,
    ConfigDict,
    PositiveInt,
    StringConstraints,
    ValidationError,
)

from owsclient.backoff import get_backoff_with_full_jitter
from owsclient.constants import PROD_AUTH0_URL, PROD_ENVIRONMENT, QA_AUTH0_URL
from owsclient.logging_adapter import get_logger
from owsclient.m2m.base import (
    DEFAULT_CACHE_KEY,
    DEFAULT_LEEWAY_SECONDS,
    BaseM2MTokenManager,
    M2MToken,
)
from owsclient.protocols import Cache, SecretsManager

logger = get_logger(__name__)

DELIMITER_ITEM = "|"
DELIMITER_KEY_VAL = "#"
AUTH0_CREDENTIALS_CACHE_KEY_PREFIX = "_auth0_client_credentials"
M2M_CLIENT_CREDENTIALS_SECRET_NAME = "M2M_AUTH0_CLIENT_CREDENTIALS"
OAUTH_TOKEN_ENDPOINT = "/oauth/token"
IMPERSONATION_UUID_FIELD = "impersonate_identity_uuid"
RETRYABLE_HTTP_STATUS_CODES = (429, 502, 503, 504)


class ClientCredentials(BaseModel):
    """Auth0 Client credentials."""

    audience: Annotated[str, StringConstraints(min_length=1)]
    client_id: Annotated[str, StringConstraints(min_length=1)]
    client_secret: Annotated[str, StringConstraints(min_length=1)]
    grant_type: Annotated[str, StringConstraints(min_length=1)] = "client_credentials"

    model_config = ConfigDict(str_strip_whitespace=True)


class OAuthToken(BaseModel):
    """OAuth token response from Auth0 APIs."""

    access_token: Annotated[str, StringConstraints(min_length=1)]
    expires_in: PositiveInt
    token_type: Annotated[str, StringConstraints(min_length=1)]

    model_config = ConfigDict(str_strip_whitespace=True)

    def convert_to_m2m_token(self) -> M2MToken:
        """Convert the auth0 response schema to an M2M pydantic schema."""
        return M2MToken.model_validate(
            {
                "token": self.access_token,
                "expires_at": datetime.datetime.now(datetime.UTC)
                + datetime.timedelta(seconds=self.expires_in),
            }
        )


class Auth0TokenException(Exception):
    """Exception class for Auth0 errors."""

    ...


def generate_jwt_token_with_impersonation_with_retries(
    client_credentials: ClientCredentials,
    impersonated_identity_uuid: str,
    auth0_url: str,
    timeout: float = 10.0,
    connect_retries: int = 3,
    server_error_retries: int = 3,
) -> OAuthToken:
    """Generate an OAuth token for impersonating a user.

    Calls Auth0's `/oauth/token` endpoint with client credentials and the impersonated identity UUID.

    Args:
    ----
        client_credentials: Auth0 client credentials.
        impersonated_identity_uuid: UUID of the identity to impersonate.
        auth0_url: Auth0 API URL
        timeout: Timeout in seconds for the HTTP request (default: 10.0).
        connect_retries: Number of times to retry when there is a ConnectError or ConnectTimeout
        server_error_retries: Number of times to retry (with backoff) when request fails for 5xx

    """
    with httpx.Client(
        base_url=auth0_url,
        transport=httpx.HTTPTransport(retries=connect_retries),
        timeout=timeout,
    ) as client:
        server_error_retry_attempts = 0
        response = _generate_jwt_token_with_impersonation(
            client_credentials=client_credentials,
            impersonated_identity_uuid=impersonated_identity_uuid,
            client=client,
        )
        while (
            response.status_code != 200
            and server_error_retry_attempts < server_error_retries
        ):
            if response.status_code not in RETRYABLE_HTTP_STATUS_CODES:
                logger.warning(
                    "ImpersonationM2MTokenManager unable to generate JWT",
                    extra={
                        "response_status_code": response.status_code,
                        "impersonated_identity_uuid": impersonated_identity_uuid,
                        "client_credential_secret_name": client_credentials.client_id,
                        "auth0_url": auth0_url,
                        "attempt": server_error_retry_attempts,
                    },
                )
                raise Auth0TokenException(
                    "JWT Access Token could not be generated", response.text
                )
            logger.info(
                "ImpersonationM2MTokenManager is retrying request to generate a JWT",
                extra={
                    "response_status_code": response.status_code,
                    "impersonated_identity_uuid": impersonated_identity_uuid,
                    "client_credential_secret_name": client_credentials.client_id,
                    "auth0_url": auth0_url,
                    "attempt": server_error_retry_attempts,
                },
            )
            server_error_retry_attempts += 1
            sleep(get_backoff_with_full_jitter(1, 3, server_error_retry_attempts))
            response = _generate_jwt_token_with_impersonation(
                client_credentials=client_credentials,
                impersonated_identity_uuid=impersonated_identity_uuid,
                client=client,
            )

    # Success
    if response.status_code == 200:
        auth_token_value = response.json()
        try:
            return OAuthToken.model_validate(auth_token_value)
        except ValidationError as e:
            raise Auth0TokenException("JWT Access Token could not be parsed") from e
    # Failure
    logger.warning(
        "ImpersonationM2MTokenManager unable to generate JWT",
        extra={
            "impersonated_identity_uuid": impersonated_identity_uuid,
            "client_credential_secret_name": client_credentials.client_id,
            "auth0_url": auth0_url,
            "attempt": server_error_retry_attempts,
        },
    )
    raise Auth0TokenException(
        f"JWT Access Token could not be generated after {server_error_retries} attempts",
        response.text,
    )


def _generate_jwt_token_with_impersonation(
    client_credentials: ClientCredentials,
    impersonated_identity_uuid: str,
    client: httpx.Client,
) -> httpx.Response:
    """Generate an OAuth token for impersonating a user.

    Calls Auth0's `/oauth/token` endpoint with client credentials and the impersonated identity UUID.

    Args:
    ----
        client_credentials: Auth0 client credentials.
        impersonated_identity_uuid: UUID of the identity to impersonate.
        client: httpx.Client

    """
    payload = client_credentials.model_dump()
    payload[IMPERSONATION_UUID_FIELD] = impersonated_identity_uuid

    try:
        response = client.post(OAUTH_TOKEN_ENDPOINT, json=payload)
    except httpx.HTTPError as exc:
        raise Auth0TokenException(
            f"Network error while requesting JWT Access Token: {exc}"
        ) from exc

    return response


class ImpersonationM2MTokenManager(BaseM2MTokenManager):
    """Class to fetch Auth0 app credentials and generate an impersonation JWT."""

    def __init__(
        self,
        secrets_manager: SecretsManager,
        environment: str,
        service_name: str,
        leeway_seconds: int = DEFAULT_LEEWAY_SECONDS,
        cache: Cache | dict[str, Any] | None = None,
        cache_key: str | None = None,
        auth0_request_timeout: float = 10.0,
        auth0_url: str | None = None,
    ) -> None:
        """Create an ImpersonationM2MTokenManager instance."""
        super().__init__(
            secrets_manager=secrets_manager,
            environment=environment,
            service_name=service_name,
            leeway_seconds=leeway_seconds,
            cache_key=cache_key or DEFAULT_CACHE_KEY,
        )
        self._cache = cache or {}
        self._auth0_request_timeout = auth0_request_timeout

        # Allow the caller to override auth0 url
        # Fallback to the normal QA or PROD urls.
        self._auth0_url = auth0_url or (
            PROD_AUTH0_URL if environment == PROD_ENVIRONMENT else QA_AUTH0_URL
        )

    def generate_secret_name(self) -> str:
        """Generate the full secret name from service, env, and secret_name."""
        return f"{self.environment}/{self.service_name}/{M2M_CLIENT_CREDENTIALS_SECRET_NAME}"

    def generate_auth0_credentials_cache_key(
        self,
    ) -> str:
        """Generate a cache key for the auth0 client credentials.

        Ex: "_auth0_client_credentials|service_name#test-service"
        """
        return f"{AUTH0_CREDENTIALS_CACHE_KEY_PREFIX}{DELIMITER_ITEM}service_name{DELIMITER_KEY_VAL}{self.service_name}"

    def generate_m2m_token_cache_key(self, impersonated_identity_uuid: str) -> str:
        """Generate a cache key for an impersonation json web token (JWT).

        Ex: "_m2m_token|service_name#my-app|impersonated_identity_uuid#f94b0c5a-b520-486b-ac17-e59e9888b8bd"
        """
        keyval_joined = [
            DELIMITER_KEY_VAL.join(kv)
            for kv in [
                ("service_name", self.service_name),
                ("impersonated_identity_uuid", impersonated_identity_uuid),
            ]
        ]
        return f"{self.cache_key}{DELIMITER_ITEM}{DELIMITER_ITEM.join(keyval_joined)}"

    def _get_client_credentials_from_secret_manager(self: Self) -> ClientCredentials:
        """Use SecretsManager to get the machine's M2M_CLIENT_CREDENTIALS."""
        client_credentials_json = self._secrets_manager.get_secret(
            secret_name=self.generate_secret_name()
        )
        if isinstance(client_credentials_json, dict):
            return ClientCredentials.model_validate(client_credentials_json)
        return ClientCredentials.model_validate_json(client_credentials_json)

    def _get_client_credentials_from_cache(self: Self) -> ClientCredentials | None:
        """Get the cached the client credentials so we don't have to fetch from SecretsManager every time."""
        cache_key = self.generate_auth0_credentials_cache_key()
        credentials_json = self._cache.get(cache_key)
        if not credentials_json:
            return None

        try:
            credentials_model = ClientCredentials.model_validate_json(credentials_json)
            return credentials_model
        except ValidationError:
            logger.info(
                "Ignoring the cache. ClientCredentials is not valid",
                extra={
                    "library": {
                        "name": "python-owsclient",
                        "language": "python",
                    },
                },
            )

        return None

    def _get_client_credentials(self) -> ClientCredentials:
        """Get the client credentials from SecretsManager or cache."""
        client_credentials = self._get_client_credentials_from_cache()

        if not client_credentials:
            client_credentials = self._get_client_credentials_from_secret_manager()

            cache_key = self.generate_auth0_credentials_cache_key()
            client_credentials_str = client_credentials.model_dump_json()

            if isinstance(self._cache, Cache):
                self._cache.set(cache_key, value=client_credentials_str, timeout=0)
            else:
                self._cache[cache_key] = client_credentials_str

        return client_credentials

    def generate_auth_token(
        self: Self,
        impersonated_identity_uuid: str,
    ) -> M2MToken:
        """Use the client credentials and the impersonated_identity_uuid to get a fresh JWT (plus expiration)."""
        client_credentials = self._get_client_credentials()

        oauth_token = generate_jwt_token_with_impersonation_with_retries(
            client_credentials=client_credentials,
            impersonated_identity_uuid=impersonated_identity_uuid,
            timeout=self._auth0_request_timeout,
            auth0_url=self._auth0_url,
        )
        return oauth_token.convert_to_m2m_token()

    def _get_token_payload_from_cache(
        self: Self, impersonated_identity_uuid: str
    ) -> M2MToken | None:
        """Get the M2MToken obj from cache for the impersonated_identity_uuid."""
        cache_key = self.generate_m2m_token_cache_key(
            impersonated_identity_uuid=impersonated_identity_uuid
        )
        token = self._cache.get(cache_key)
        if not token:
            return None
        try:
            token_model = M2MToken.model_validate_json(token)
            return token_model
        except ValidationError:
            logger.info(
                "Ignoring the cache, M2MToken is not valid",
                extra={
                    "library": {
                        "name": "python-owsclient",
                        "language": "python",
                    },
                },
            )
        return None

    def get_token_string(self: Self, impersonated_identity_uuid: str) -> str:
        """Return the impersonated JWT token."""
        m2m_token = self._get_token_payload_from_cache(
            impersonated_identity_uuid=impersonated_identity_uuid,
        )

        if m2m_token and not m2m_token.is_expired(leeway_seconds=self.leeway_seconds):
            return m2m_token.token

        m2m_token = self.generate_auth_token(
            impersonated_identity_uuid=impersonated_identity_uuid
        )

        m2m_token_str = m2m_token.model_dump_json()
        cache_key = self.generate_m2m_token_cache_key(
            impersonated_identity_uuid=impersonated_identity_uuid
        )

        if isinstance(self._cache, Cache):
            self._cache.set(cache_key, value=m2m_token_str, timeout=0)
        else:
            self._cache[cache_key] = m2m_token_str

        return m2m_token.token