from __future__ import annotations

import json
from typing import Any

from django.core.serializers.json import DjangoJSONEncoder
from django.db import models
from django.db.backends.base.base import BaseDatabaseWrapper
from django.db.models.expressions import Expression

from saas.utils.secure import decrypt_data, encrypt_data


class BaseEncryptedField(models.TextField):
    def get_db_prep_value(
        self,
        value: Any,
        connection: BaseDatabaseWrapper,  # noqa: ARG002
        prepared: bool = False,  # noqa: FBT001, FBT002
    ) -> Any:
        if not prepared:
            value = self.get_prep_value(value)
        return value

    def from_db_value(
        self,
        value: Any,
        expression: Expression,  # noqa: ARG002
        connection: BaseDatabaseWrapper,  # noqa: ARG002
    ) -> Any:
        return self.to_python(value)

    def clean(self, value: Any, model_instance: models.Model) -> Any:
        """
        Create and assign a semaphore so that to_python method will not try
        to decrypt an already decrypted value during cleaning of a form
        """
        setattr(self, '_already_decrypted', True)
        try:
            return super().clean(value, model_instance)
        finally:
            if hasattr(self, '_already_decrypted'):
                delattr(self, '_already_decrypted')


class EncryptedBinaryField(BaseEncryptedField):
    description = 'A field that stores bytes encrypted with JWE (X25519)'

    def get_prep_value(self, value: bytes | None) -> str | None:
        """Encrypt Python bytes into a JWE string for the database."""
        if value is None:
            return None

        if not isinstance(value, (bytes, memoryview)):
            raise TypeError(f'EncryptedBinaryField must receive bytes, got {type(value)}')

        return encrypt_data(bytes(value))

    def to_python(self, value: Any) -> bytes | None:
        """Decrypt a JWE string from the database back into Python bytes with Key Rotation support."""
        if value is None or hasattr(self, '_already_decrypted'):
            return value

        if isinstance(value, str):
            return decrypt_data(value)

        return value


class EncryptedTextField(BaseEncryptedField):
    description = 'A field that stores text encrypted with JWE (X25519)'

    def get_prep_value(self, value: str | None) -> str | None:
        """Encrypt Python string into a JWE string for the database."""
        if value is None:
            return None

        if not isinstance(value, str):
            raise TypeError(f'EncryptedTextField must receive str, got {type(value)}')

        return encrypt_data(value.encode('utf-8'))

    def to_python(self, value: Any) -> str | None:
        """Decrypt a JWE string from the database back into Python string with Key Rotation support."""
        if value is None or hasattr(self, '_already_decrypted'):
            return value

        if isinstance(value, str):
            decrypted = decrypt_data(value)
            if decrypted is not None:
                return decrypted.decode('utf-8')
            return None

        return value


class EncryptedJSONField(BaseEncryptedField):
    description = 'A field that stores JSON encrypted with JWE (X25519)'

    def get_prep_value(self, value: Any | None) -> str | None:
        """Encrypt Python object into a JWE string for the database."""
        if value is None:
            return None

        return encrypt_data(json.dumps(value, cls=DjangoJSONEncoder).encode('utf-8'))

    def to_python(self, value: Any) -> Any:
        """Decrypt a JWE string from the database back into Python object with Key Rotation support."""
        if value is None or hasattr(self, '_already_decrypted'):
            return value

        if isinstance(value, str):
            decrypted = decrypt_data(value)
            if decrypted is not None:
                return json.loads(decrypted.decode('utf-8'))
            return None

        return value
