import base64
from lclazy import LazyLoader
from aspen_crypto.keys import (
    KeyDetails,
    generate_kid,
    load_local_umbrella_api_key,
    PublicKey,
)
from datetime import datetime, timedelta, timezone
from sequential_uuids.generators import uuid_time_nextval
from uuid import UUID

from django.conf import settings
from .models import EncryptionKey
from . import get_setting
from .typehints import CidType, KeyData

pyseto = LazyLoader('pyseto', globals(), 'pyseto')
_retrieve_key_from_aspen = None


# ---------------------------------------------------------------------------
class KeyLoaderBase:
    """
    Public keys can be loaded from:
       * config file
       * database
       * request body
       * Aspen API call

    This class implements the basic interface for such loaders.
    """

    def __init__(self, cid: CidType = '', kid: str = '', data: dict | None = None) -> None:  # type: ignore[type-arg]
        self.key_cid = cid
        self.key_kid = kid
        self.data = data
        self.key_data: KeyData | KeyDetails | None = None

    def update_last_used(self) -> None:
        """
        Implement on subclasses. Mostly going to be
        used for keys retrieved from the database.
        """
        pass

    def public_key(self):
        """
        Implement on subclasses.
        """
        raise NotImplementedError()


class DatabaseLoader(KeyLoaderBase):
    """
    Load Key from EncryptionKey record.
    """

    def public_key(self) -> KeyData | None:
        """
        Get the Key ID or Canonical ID of the key that signed this token.
        """
        if self.key_cid:
            query_params = {'cid': self.key_cid}
        elif self.key_kid:
            query_params = {'kid': self.key_kid}
        else:
            return

        key_data = (
            EncryptionKey.objects.filter(**query_params)
            .values_list(
                'cid',
                'kid',
                'key',
                'last_used_on',
                'status',
                'subject',
                'revocation_date',
                named=True,
            )
            .first()
        )

        if key_data:
            if key_data.status != EncryptionKey.STATUS.active:
                raise pyseto.VerifyError('Key has been revoked')

            self.key_data = KeyData(
                cid=key_data.cid,
                kid=key_data.kid,
                key=key_data.key,
                subject=key_data.subject,
                last_used_on=key_data.last_used_on,
            )

        return self.key_data

    def update_last_used(self) -> None:
        """
        Update last_used_on date if it hasn't been updated recently.
        Not updating every single time, since a token might be
        used many times per minute.
        """
        if not self.key_data or isinstance(self.key_data, KeyDetails):
            return

        last_used_on = self.key_data.last_used_on
        now = datetime.now(tz=timezone.utc)

        last_used_delta = timedelta(
            minutes=get_setting('LAST_USED_INTERVAL')  # zuban: ignore[arg-type]
        )
        if not last_used_on or now - last_used_on > last_used_delta:
            EncryptionKey.objects.filter(cid=self.key_data.cid).update(last_used_on=now)


class AspenLoader(KeyLoaderBase):
    """
    Load Key from Aspen API call.
    """

    def public_key(self) -> KeyData | None:
        if not settings.CHECK_ASPEN_FOR_UNKNOWN_KEYS:
            return

        global _retrieve_key_from_aspen
        if _retrieve_key_from_aspen is None:
            from .aspen import retrieve_key_from_aspen

            _retrieve_key_from_aspen = retrieve_key_from_aspen

        try:
            # Calls to Aspen will fail if there's no local key to load.
            self.key_data = _retrieve_key_from_aspen(self.key_cid, self.key_kid)
        except ValueError:
            return

        return self.key_data


class LocalKeyLoader(KeyLoaderBase):
    """
    Load Key from local configuration file.
    """

    def public_key(self) -> KeyDetails | None:
        """
        On systems with a local key configured, this
        loader will always have a value. But the Token
        must provide a matching `cid` or `kid` value.
        """
        try:
            key_data = load_local_umbrella_api_key()
        except ValueError:
            return

        if key_data.cid == self.key_cid or key_data.kid == self.key_kid:
            self.key_data = key_data
            return key_data


class PemLoader(KeyLoaderBase):
    """
    Load Key from the specified PEM.

    On limited occasions, such as registrations, when
    a registration secret is otherwise provided, a
    token can be signed with the PEM that's provided
    in the request body.
    """

    def public_key(self) -> KeyDetails | None:
        if not self.data or not (pem := self.data.get('pem')):
            return

        if isinstance(pem, str):
            pem = base64.b64decode(pem) if pem.endswith('=') else pem.encode()

        error, pub_key = PublicKey.load_serialized_public_key(pem)
        if error or not pub_key:
            return

        if (cid := self.data.get('cid')) and isinstance(cid, str):
            cid = UUID(cid)

        key_data = KeyDetails(
            cid=cid or uuid_time_nextval(),
            kid=generate_kid(pub_key.as_pem),
            key=None,
            public_key=pub_key,
        )

        self.key_data = key_data
        return key_data
