import base64
import os
from datetime import UTC, datetime
from typing import Protocol

import jwt
from cryptography.hazmat.primitives import hashes, padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC


class Encrypter(Protocol):
    def encrypt(self, payload: bytes) -> bytes: ...

    def decrypt(self, token: bytes | str) -> bytes: ...


class RawEncrypter(Encrypter):
    def encrypt(self, payload: bytes) -> bytes:
        return base64.b64encode(payload)

    def decrypt(self, token: bytes | str) -> bytes:
        if isinstance(token, str):
            token = token.encode()
        return base64.b64decode(token)


class JWTEncrypter(Encrypter):
    AES256_BLOCK_SIZE = 128

    def __init__(self, *, key: str) -> None:
        signing_key, password, salt = key.split(" ")
        self._signing_key = signing_key.encode("utf-8")
        self._encryption_key = self._derive_encryption_key(
            password=password.encode("utf-8"),
            salt=bytes.fromhex(salt),
        )

    @classmethod
    def generate_key(cls) -> str:
        signing_key = base64.b85encode(os.urandom(32)).decode("utf-8")
        encryption_password = base64.b85encode(os.urandom(24)).decode("utf-8")
        encryption_salt = os.urandom(8).hex()
        return " ".join([signing_key, encryption_password, encryption_salt])

    @classmethod
    def _derive_encryption_key(cls, password: bytes, salt: bytes) -> bytes:
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA1(),
            length=32,
            salt=salt,
            iterations=1000,
        )
        return kdf.derive(password)

    @staticmethod
    def _generate_iv() -> str:
        return os.urandom(16).hex()

    def _get_cipher(self, iv: bytes) -> Cipher[modes.CBC]:
        return Cipher(
            algorithm=algorithms.AES256(self._encryption_key),
            mode=modes.CBC(iv),
        )

    def encrypt(self, payload: bytes) -> bytes:
        iv = self._generate_iv()
        ciphertext = self._encrypt_payload(payload, iv=bytes.fromhex(iv))
        encrypted_payload = iv + base64.b64encode(ciphertext).decode("utf-8")
        return jwt.encode(
            {
                "payload": encrypted_payload,
                "iat": int(datetime.now(tz=UTC).timestamp()),
            },
            key=self._signing_key,
            algorithm="HS256",
        ).encode()

    def _encrypt_payload(self, payload: bytes, *, iv: bytes) -> bytes:
        padder = padding.PKCS7(block_size=self.AES256_BLOCK_SIZE).padder()
        padded_data = padder.update(payload) + padder.finalize()
        encryptor = self._get_cipher(iv).encryptor()
        return encryptor.update(padded_data) + encryptor.finalize()

    def decrypt(self, token: bytes | str) -> bytes:
        jwt_payload = jwt.decode(token, key=self._signing_key, algorithms=["HS256"])
        iv, encrypted_payload_b64 = (
            jwt_payload["payload"][:32],
            jwt_payload["payload"][32:],
        )
        ciphertext = base64.b64decode(encrypted_payload_b64)
        return self._decrypt_payload(ciphertext, iv=bytes.fromhex(iv))

    def _decrypt_payload(self, ciphertext: bytes, *, iv: bytes) -> bytes:
        decryptor = self._get_cipher(iv=iv).decryptor()
        plaintext_padded = decryptor.update(ciphertext)
        plaintext_padded += decryptor.finalize()
        unpadder = padding.PKCS7(self.AES256_BLOCK_SIZE).unpadder()
        return unpadder.update(plaintext_padded) + unpadder.finalize()