"""M2M Token Managers."""

from __future__ import annotations

import datetime
import json
from typing import Any, Self

from anyio import to_thread
from pydantic import (
    AwareDatetime,
    BaseModel,
    ConfigDict,
    ValidationError,
    field_validator,
)

from owsclient.logging_adapter import get_logger
from owsclient.protocols import AsyncCache, Cache, SecretsManager

EXPIRATION_FORMAT = "%Y-%m-%d %H:%M:%S.%f%z"
DEFAULT_LEEWAY_SECONDS = 60
M2M_JWT_ACCESS_TOKEN_SECRET_NAME = "M2M_JWT_ACCESS_TOKEN"
DEFAULT_CACHE_KEY = "_m2m_token"

logger = get_logger(__name__)


class M2MToken(BaseModel):
    """Representation of an M2M Token."""

    token: str
    expires_at: AwareDatetime

    model_config = ConfigDict(
        str_strip_whitespace=True,
    )

    @field_validator("token")
    @classmethod
    def check_action_not_empty(cls, v: Any) -> Any:
        """Validate token is not empty."""
        assert v != "", "Empty strings are not allowed."
        return v

    def is_expired(self, leeway_seconds: int | None = None) -> bool:
        """Return True if the M2M Token is expired."""
        return self.expires_at <= (
            datetime.datetime.now(datetime.UTC)
            + datetime.timedelta(seconds=leeway_seconds or 0)
        )


class BaseM2MTokenManager:
    """BaseM2MTokenManager class."""

    def __init__(
        self,
        secrets_manager: SecretsManager,
        environment: str,
        service_name: str,
        leeway_seconds: int = DEFAULT_LEEWAY_SECONDS,
        cache_key: str | None = DEFAULT_CACHE_KEY,
    ) -> None:
        """Create a BaseM2MTokenManager instance."""
        self.environment = environment
        self.service_name = service_name
        self.leeway_seconds = leeway_seconds
        self._secrets_manager = secrets_manager
        self.cache_key = cache_key or DEFAULT_CACHE_KEY

    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_JWT_ACCESS_TOKEN_SECRET_NAME}"
        )


class M2MTokenManager(BaseM2MTokenManager):
    """Class to fetch and extract bearer token secrets."""

    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,
    ) -> None:
        """Create a M2MTokenManager 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 {}

    def _get_token_payload_from_secret_manager(self: Self) -> str:
        """Fetch the token secret using the _secret_manager."""
        m2m_jwt_info_json = self._secrets_manager.get_secret(
            secret_name=self.generate_secret_name()
        )
        if isinstance(m2m_jwt_info_json, dict):
            return json.dumps(m2m_jwt_info_json)
        return m2m_jwt_info_json

    def _get_token_from_cache(self) -> M2MToken | None:
        """Fetch the token from cache."""
        token = self._cache.get(self.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) -> str:
        """Fetch the secret value and extract the `token`."""
        token = self._get_token_from_cache()
        if token and not token.is_expired(leeway_seconds=self.leeway_seconds):
            return token.token

        token_str = self._get_token_payload_from_secret_manager()
        token = M2MToken.model_validate_json(token_str)

        if token.is_expired():
            logger.warning(
                "Fetched an expired token from secrets manager. Please manually rotate '%s'",
                self.generate_secret_name(),
                extra={
                    "library": {
                        "name": "python-owsclient",
                        "language": "python",
                    },
                },
            )
            # Don't cache an expired token. Return immediately.
            return token.token

        # Cache token as string forever and always rely on is_expired() check
        if isinstance(self._cache, Cache):
            self._cache.set(self.cache_key, value=token_str, timeout=0)
        else:
            self._cache[self.cache_key] = token_str

        return token.token


class AsyncM2MTokenManager(BaseM2MTokenManager):
    """Class to fetch and extract bearer token secrets."""

    def __init__(
        self,
        secrets_manager: SecretsManager,
        environment: str,
        service_name: str,
        leeway_seconds: int = DEFAULT_LEEWAY_SECONDS,
        cache: AsyncCache | dict[str, Any] | None = None,
        cache_key: str | None = None,
    ) -> None:
        """Create a AsyncM2MTokenManager 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 {}

    async def _get_token_payload_from_secret_manager(
        self: Self,
    ) -> str:
        """Fetch the token secret using the _secret_manager."""
        m2m_jwt_info_json = await to_thread.run_sync(
            self._secrets_manager.get_secret,
            self.generate_secret_name(),
        )
        if isinstance(m2m_jwt_info_json, dict):
            return json.dumps(m2m_jwt_info_json)
        return m2m_jwt_info_json

    async def _get_token_from_cache(self) -> M2MToken | None:
        """Fetch the token from cache."""
        if isinstance(self._cache, AsyncCache):
            token = await self._cache.get(self.cache_key)
        else:
            token = self._cache.get(self.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

    async def get_token_string(self) -> str:
        """Fetch the secret value and extract the `token`."""
        token = await self._get_token_from_cache()
        if token and not token.is_expired(leeway_seconds=self.leeway_seconds):
            return token.token

        token_str = await self._get_token_payload_from_secret_manager()
        token = M2MToken.model_validate_json(token_str)

        if token.is_expired():
            logger.warning(
                "Fetched an expired token from secrets manager. Please manually rotate '%s'",
                self.generate_secret_name(),
                extra={
                    "library": {
                        "name": "python-owsclient",
                        "language": "python",
                    },
                },
            )

            # Don't cache an expired token. Return immediately.
            return token.token

        # Cache token as string forever and always rely on is_expired() check
        if isinstance(self._cache, AsyncCache):
            await self._cache.set(self.cache_key, value=token_str, ttl=0)
        else:
            self._cache[self.cache_key] = token_str

        return token.token