"""
Enhanced key management with rotation support.
"""

import asyncio
import base64
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
import enum
from logging import Logger
from uuid import uuid4

from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey
from cryptography.hazmat.primitives.asymmetric.types import (
    PrivateKeyTypes,
    PublicKeyTypes,
)

from phederation.cache import BaseCache
from phederation.cache.base import WithCache, with_cache
from phederation.federation.resolver import ActivityPubResolver
from phederation.models import APPublicKey, DataIntegrityProof
from phederation.models.keys import APPrivateKey
from phederation.storage.base import StorageBackend
from phederation.utils import ObjectId
from phederation.utils.base import AccessType, UrlType, assemble_id_url
from phederation.utils.exceptions import (
    KeyManagementError,
    ResolverError,
    SecurityError,
    catch_exceptions,
)
from phederation.utils.logging import configure_logger
from phederation.utils.settings import PhedSettings


class KeyType(str, enum.Enum):
    RSA = "rsa-sha256-key"
    Fernet = "fernet-key"


@dataclass
class KeyPair:
    key_id: ObjectId
    private_key: RSAPrivateKey | None = None
    public_key: RSAPublicKey | None = None
    created_at: datetime | None = None
    expires_at: datetime | None = None


class KeyManager(WithCache):
    """
    Enhanced key management with rotation.

    This class handles private/public RSA key generation.
    Keys are stored on disk, and stored in an "active_keys" dictionary for easy access.
    Keys expire after a set amount of days (specified in the KeyRotation settings).
    The key_id is an url that points to the location where information about the public key can be gathered.

    Note that keys are not associated to any user or the instance, this is handled outside of the class.
    """

    def __init__(
        self,
        settings: PhedSettings,
        resolver: ActivityPubResolver,
        storage: StorageBackend,
        cache: BaseCache | None = None,
    ):
        """
        Creates a new KeyManager.

        Args:
            settings (PhedSettings): to read the domain url of the current instance, and key rotation settings.
            resolver (ActivityPubResolver): resolver object to resolve keys over the network.
            storage (StorageBackend): storage object to store and retrieve keys in a database.
            cache (BaseCache | None, optional): Cache to store remote public keys. Defaults to None.
        """
        self.settings: PhedSettings = settings
        self.resolver: ActivityPubResolver = resolver
        self.storage: StorageBackend = storage
        self.cache: BaseCache | None = cache
        self.logger: Logger = configure_logger(__name__, prefix=self.settings.federation.logging_prefix)

    async def initialize(self) -> None:
        """Initialize key manager."""
        pass

    async def resolve_current_key(self, key_ids: list[ObjectId] | ObjectId | None):
        """From the given list of key ids, return the most recently created one.

        This is useful to get the most recent key of an APActor, especially after key rotation.

        Args:
            key_ids (list[ObjectId]): List of key ids.
        """
        if not key_ids:
            return None
        if isinstance(key_ids, str):
            key_ids = [key_ids]

        keys = await asyncio.gather(*[self.resolve_key(key_id) for key_id in key_ids])

        def sort_key_by_creation_date(key: APPublicKey):
            return key.created_at or datetime.now(timezone.utc)

        keys.sort(key=sort_key_by_creation_date, reverse=True)
        if len(keys) > 0:
            return keys[0]
        else:
            return None

    @with_cache(prefix=__name__ + "get_key_id_from_actor")
    async def get_key_id_from_actor(self, actor_id: ObjectId) -> ObjectId:
        """Obtain the valid id of the given actor.

        Args:
            actor_id (ObjectId): id of the actor

        Raises:
            ResolverError: If the actor or the publicKey property could not be resolved.

        Returns:
            str: the local id (not the URL, but the integer!) of the key of the actor.
        """
        actor = await self.resolver.resolve_actor(actor_id=actor_id)
        if not actor:
            self.logger.warning(f"Failed to read actor with id {actor_id}")
            raise ResolverError(f"Could not read actor")
        all_keys = actor.public_key_as_ids()
        self.logger.debug(f"Resolved actor {actor_id}, get current key from keys {all_keys}...")
        public_key = await self.resolve_current_key(all_keys)

        if public_key and public_key.id:
            return public_key.id
        else:
            self.logger.warning(f"Failed to get 'public_key' property from actor with id {actor_id}; dict={actor.serialize()}")
            raise ResolverError(f"Could not resolve actor property 'public_key / publicKey'. Maybe not a local actor")

    @catch_exceptions(KeyManagementError, "Fernet key generation failed")
    async def generate_fernet_key(self, actor_id: ObjectId | None) -> tuple[ObjectId, bytes]:
        """Generates a new Fernet key and stores it to storage.

        Note that this key is stored unencrypted in the storage database.
        Anyone with direct access to that database can access all information that is encrypted with this key.

        Usually only the messages to and from the temporalio backend are encrypted with this key, so
        this should not be an issue.

        Args:
            actor_id: The id of the actor this key is associated with.

        Returns:
            ObjectId, bytes: tuple of (The key id, the raw Fernet key).
        """
        # generate the key
        key_bytes = Fernet.generate_key()

        public_pem = "NO PUBLIC KEY PEM"
        private_pem = base64.encodebytes(key_bytes).decode()

        # Set validity period
        created_at = datetime.now(timezone.utc)
        expires_at = created_at + timedelta(seconds=self.settings.security.key_rotation_period)

        # Generate key ID (for HTTP use)
        key_primary = str(uuid4())
        key_id = assemble_id_url(type=UrlType.Keys, base_url=self.settings.domain.base_url, primary=key_primary)

        key = APPrivateKey(
            id=key_id,
            attributed_to=actor_id,
            media_type=KeyType.Fernet.value,
            public_key_pem=public_pem,
            private_key_pem=private_pem,
            created_at=created_at,
            expires_at=expires_at,
            visibility=AccessType.PRIVATE.value,
        )
        # TODO: also add this key to the list of keys of the given actor, so that it can be replaced during key rotation
        key_id = await self.storage.key.create(data=key)
        return key_id, key_bytes

    @catch_exceptions(KeyManagementError, "RSA key generation failed")
    async def generate_key_pair(self, actor_id: ObjectId | None = None) -> KeyPair:
        """Generates a new RSA key pair and stores it to storage.

        Parameters:
            actor_id (ObjectId): if the key is for an actor, provide the id here. Default: None

        Raises:
            KeyManagementError: If the key generation fails.

        Returns:
            KeyPair: The generated keys (private and public key).
        """
        self.logger.debug("Generating new key pair")
        # Generate keys
        private_key = rsa.generate_private_key(public_exponent=65537, key_size=self.settings.security.key_size)
        public_key = private_key.public_key()

        # Set validity period
        created_at = datetime.now(timezone.utc)
        expires_at = created_at + timedelta(seconds=self.settings.security.key_rotation_period)

        # Generate key ID (for HTTP use)
        key_primary = str(uuid4())
        key_id = assemble_id_url(type=UrlType.Keys, base_url=self.settings.domain.base_url, primary=key_primary)

        self.logger.debug(f"Generated key ID: {key_id}")

        # Create key pair
        key_pair = KeyPair(private_key=private_key, public_key=public_key, created_at=created_at, expires_at=expires_at, key_id=key_id)

        # Save keys in storage
        await self._save_RSA_key_pair(key_pair, actor_id=actor_id)
        self.logger.debug("Saved key pair to storage")

        return key_pair

    @with_cache(__name__ + "_resolve_key")
    async def resolve_key(self, key_id: ObjectId) -> APPrivateKey | APPublicKey:
        """Get the APPrivateKey for a given key_id.
        If only the public key is available, return an APPRivateKey or APPublicKey.

        Args:
            key_id (str): id number of the given key

        Returns:
            APPrivateKey | APPublicKey: The key data associated to the id.
        """
        # check if the key is available locally
        key_obj = await self.storage.key.read(id=key_id)

        # if not, resolve over network
        if not key_obj or not key_obj.public_key_pem:
            if self.resolver.is_local(url=key_id):
                raise ResolverError(f"Key {key_id} could not be found locally, but has a local url")
            self.logger.debug(f"Key {key_id} is not local, resolving public key over network")
            # given an actor, check if we have its public key already. If not, resolve it and add it to the key manager.
            key_obj = await self.resolver.resolve_object(object_id=key_id)
            if isinstance(key_obj, APPrivateKey) and key_obj.public_key_pem:
                self.logger.debug(f"Key loaded successfully: keyId={key_id}")
            else:
                self.logger.debug(f"KeyManager: object_dict={key_obj}, not resolved for key with id={key_id}.")
                raise ResolverError(
                    f"KeyManager: object_dict={key_obj}, not resolved for key with id={key_id}.", user_facing_message="Could not resolve key"
                )

        if not key_obj.private_key_pem:
            key_obj = APPublicKey.deserialize(key_obj.serialize())
            if not isinstance(key_obj, APPublicKey):
                raise ResolverError(f"Could not resolve public key.")
        return key_obj

    @with_cache(__name__ + "get_RSA_key_pair")
    async def get_RSA_key_pair(self, key_id: ObjectId) -> KeyPair:
        """Get the RSA KeyPair for a given key_id, including the private key.

        Args:
            key_id (str): id number of the given key

        Returns:
            KeyPair: The KeyPair associated to the id.
        """
        key = await self.resolve_key(key_id=key_id)

        public_key: PublicKeyTypes = serialization.load_pem_public_key(str(key.public_key_pem).encode())
        private_key: PrivateKeyTypes | None = None
        if isinstance(key, APPrivateKey) and key.private_key_pem:
            private_key = serialization.load_pem_private_key(str(key.private_key_pem).encode(), password=None)

        if not isinstance(public_key, RSAPublicKey):
            raise KeyManagementError("public key PEM did not resolve to RSAPublicKey")
        if private_key and not isinstance(private_key, RSAPrivateKey):
            raise KeyManagementError("private key PEM did not resolve to RSAPublicKey")

        key_pair = KeyPair(key_id=key_id, private_key=private_key, public_key=public_key, created_at=key.created_at, expires_at=key.expires_at)

        return key_pair

    async def get_public_key_APObject(self, key_id: ObjectId) -> APPublicKey:
        key = await self.resolve_key(key_id)
        if isinstance(key, APPrivateKey):
            key.private_key_pem = None
            public_key = APPublicKey.deserialize(key.serialize())
            if not isinstance(public_key, APPublicKey):
                raise KeyManagementError(f"Could not deserialize public key from private key, key_id={key_id}")
        else:
            public_key = key
        return public_key

    async def get_public_key(self, key_id: ObjectId) -> RSAPublicKey | None:
        key = await self.get_RSA_key_pair(key_id=key_id)
        return key.public_key

    async def get_public_key_pem(self, key_id: ObjectId) -> str | None:
        """Get the public key in PEM format for a key_id.

        Args:
            key_id (str): id number of the given key

        Returns:
            str: A string representation of the public PEM key file.
        """
        key = await self.get_public_key(key_id=key_id)
        if key:
            return key.public_bytes(encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo).decode("utf-8")

    async def verify_key(self, key_id: ObjectId) -> bool:
        """Verify a key's validity.
        Checks if the given key is in the active keys and not expired yet.
        Does not check external keys yet (domain is not used).

        Args:
            key_id (str): key id (URL) for the key to verify.

        Returns:
            bool: True (valid key) or False (invalid).
        """
        try:
            # Check if key is one of our active keys
            key = await self.storage.key.read(id=key_id)
            if key and key.expires_at:
                return datetime.now(timezone.utc) <= key.expires_at

            # TODO: For external keys, verify with their server implementation for external key verification
            return False

        except Exception as e:
            self.logger.error(f"Key verification failed: {e}")
            return False

    async def rotate_key(self, key: APPublicKey | ObjectId, actor_id: ObjectId | None) -> ObjectId:
        """Rotates the given key.
        If the actor id is set, also replaces the public key id for this actor with the new key.

        Args:
            key (APPublicKey): Key to replace.
            actor_id (ObjectId | None): If given, will check if the key is associated to this actor and also replace it there.

        Returns:
            ObjectId: the id of the new key, or none if it could not be generated.
        """
        # resolve key if only given an id
        if isinstance(key, str):
            key = await self.resolve_key(key_id=key)

        key_id_old = key.id
        if not key_id_old:
            raise KeyManagementError("Key id of old key before rotation should not be None")

        # replace key in actor if available
        if actor_id:
            key_id_from_actor = await self.get_key_id_from_actor(actor_id=actor_id)
            if key_id_from_actor == key.id:
                # replace the public key of the actor
                actor = await self.storage.actor.read(id=actor_id)
                if actor:
                    # generate new key
                    if key.media_type == KeyType.Fernet.value:
                        key_id_new, _ = await self.generate_fernet_key(actor_id=actor_id)
                    else:
                        key_pair = await self.generate_key_pair(actor_id=actor_id)
                        key_id_new = key_pair.key_id
                    # add new key to actor key list and update
                    actor.add_public_key(key_id_new)
                    _ = await self.storage.actor.upsert(id=actor_id, data=actor)
                    return key_id_new
                else:
                    raise KeyManagementError(f"Failed to read actor {actor_id} to replace the expired public key")
            else:
                raise KeyManagementError(f"Key id from {actor_id} is not the same as the one that should be replaced")
        else:
            raise KeyManagementError("Actor id for key is None")

    @catch_exceptions(KeyManagementError, "Failed to save key pair")
    async def _save_RSA_key_pair(self, key_pair: KeyPair, actor_id: ObjectId | None = None) -> None:
        """Save key pair to storage."""
        if not key_pair.private_key or not key_pair.public_key:
            raise KeyManagementError("Public or private key missing")

        private_pem = key_pair.private_key.private_bytes(
            encoding=serialization.Encoding.PEM, format=serialization.PrivateFormat.PKCS8, encryption_algorithm=serialization.NoEncryption()
        ).decode(encoding="utf-8")

        public_pem = key_pair.public_key.public_bytes(
            encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo
        ).decode(encoding="utf-8")

        key = APPrivateKey(
            id=key_pair.key_id,
            attributed_to=actor_id,
            public_key_pem=public_pem,
            private_key_pem=private_pem,
            media_type=KeyType.RSA.value,
            created_at=key_pair.created_at,
            expires_at=key_pair.expires_at,
        )
        await self.store_key(key)
    
    async def store_key(self, key: APPrivateKey | APPublicKey):
        if not isinstance(key, APPrivateKey):
            key = key = APPrivateKey(
                id=key.id,
                attributed_to=key.attributed_to,
                public_key_pem=key.public_key_pem,
                private_key_pem=None,
                media_type=key.media_type or KeyType.RSA.value,
                created_at=key.created_at,
                expires_at=key.expires_at,
            )
        if key.id and not await self.storage.key.read(id=key.id):
            _ = await self.storage.key.create(data=key)

    async def proof_string(self, string_to_sign: str, private_key_id: ObjectId | None) -> DataIntegrityProof:
        """Creates a DataIntegrityProof for the given string, using the private key with the given id.

        The cryptosuite is fixed to 'RSA', because the method can only deal with RSA private keys for now.
        The proof_purpose is always 'assertionMethod'.

        Args:
            string_to_sign (str): The string to create the proof for.
            private_key_id (ObjectId | None): Id of the private key to use for signing.

        Raises:
            SecurityError: If the private key cannot be found.

        Returns:
            DataIntegrityProof: An RSA proof for the given string.
        """
        if not private_key_id:
            raise SecurityError(f"Key with id {private_key_id} not found")
        proof = DataIntegrityProof(
            type="DataIntegrityProof",
            cryptosuite="RSA",
            proof_purpose="assertionMethod",
            verification_method=private_key_id,
        )
        key_pair = await self.get_RSA_key_pair(key_id=private_key_id)
        if key_pair.private_key is None:
            raise SecurityError(f"Private key {private_key_id} not found")

        proof = proof.sign(string_to_sign, key=key_pair.private_key)
        proof_id = await self.storage.proof.create(data=proof)
        proof.id = proof_id
        return proof

    async def close(self) -> None:
        """Clean up resources."""
        pass
