"""
worker_encryption.py
"""

from collections.abc import Iterable
from logging import Logger
from typing import override

from cryptography.fernet import Fernet
from temporalio.api.common.v1 import Payload
from temporalio.converter import PayloadCodec

from phederation.utils.exceptions import SecurityError
from phederation.utils.logging import configure_logger


class EncryptionCodec(PayloadCodec):
    """Handles compression and encryption of payloads to temporalio and back.

    Note that only the temporalio backend does not see the private keys,
    methods starting activities/workflows and workers handling them all run on the local instance and have access to everything.

    Encryption is done using the given key_id (typically the instance actor key), and RSA encryption with PKCS1v15 padding.

    See: https://docs.temporal.io/develop/python/converters-and-encryption
    """

    def __init__(self, fernet_key: bytes, logging_prefix: str) -> None:
        super().__init__()
        self.key: bytes = fernet_key
        self.logger: Logger = configure_logger(__name__, logging_prefix)

    @override
    async def encode(self, payloads: Iterable[Payload]) -> list[Payload]:
        return [
            Payload(
                metadata={
                    "encoding": b"binary/encrypted",
                    "encryption-key-id": b"fernet",
                },
                data=await self.encrypt(data=p.SerializeToString()),
            )
            for p in payloads
        ]

    @override
    async def decode(self, payloads: Iterable[Payload]) -> list[Payload]:
        ret: list[Payload] = []
        for p in payloads:  # Ignore ones w/out our expected encoding
            if p.metadata.get("encoding", b"").decode() != "binary/encrypted":
                ret.append(p)
                continue
            # Confirm our key ID is the same
            key_id = p.metadata.get("encryption-key-id", b"").decode(encoding="utf-8")
            if key_id != "fernet":
                raise SecurityError(f"Unrecognized key ID {key_id}. Current key ID is fernet.")
            # Decrypt and append
            ret.append(Payload.FromString(await self.decrypt(data=p.data)))
        return ret

    async def encrypt(self, data: bytes) -> bytes:
        try:
            return Fernet(self.key).encrypt(data)
        except Exception as e:
            self.logger.error(f"Error in encrypting for temporalio: {e}")
            raise SecurityError("Failed to encrypt message to temporalio")

    async def decrypt(self, data: bytes) -> bytes:
        try:
            return Fernet(self.key).decrypt(data)
        except Exception as e:
            self.logger.error(f"Error in decrypting from temporalio: {e}")
            raise SecurityError("Failed to decrypt message from temporalio")
