import base64
from gmssl.sm4 import CryptSM4, SM4_ENCRYPT, SM4_DECRYPT

from django.db import models
from django.conf import settings
from django.core.exceptions import ImproperlyConfigured


def _pkcs7_pad(data, block_size=16):
    pad_len = block_size - (len(data) % block_size)
    return data + bytes([pad_len]) * pad_len


def _pkcs7_unpad(data, block_size=16):
    if not data or len(data) % block_size != 0:
        raise ValueError('invalid data length for PKCS7 unpad')
    pad_len = data[-1]
    if pad_len < 1 or pad_len > block_size:
        raise ValueError('invalid pad length')
    if data[-pad_len:] != bytes([pad_len]) * pad_len:
        raise ValueError('invalid padding bytes')
    return data[:-pad_len]


class SM4EncryptedField(models.TextField):
    def __init__(self, *args, **kwargs):
        key = getattr(settings, 'SM4_KEY', None)
        iv = getattr(settings, 'SM4_IV', None)
        if not key or not iv:
            raise ImproperlyConfigured('SM4_KEY and SM4_IV must be configured in settings')
        self.key = key.encode('utf-8') if isinstance(key, str) else key
        self.iv = iv.encode('utf-8') if isinstance(iv, str) else iv
        if len(self.key) != 16:
            raise ImproperlyConfigured('SM4_KEY must be 16 bytes (128-bit) long')
        if len(self.iv) != 16:
            raise ImproperlyConfigured('SM4_IV must be 16 bytes long')
        super().__init__(*args, **kwargs)

    def from_db_value(self, value, expression, connection):
        if value is None:
            return value
        try:
            raw = base64.b64decode(value)
            crypt_sm4 = CryptSM4()
            crypt_sm4.set_key(self.key, SM4_DECRYPT)
            decrypted_bytes = crypt_sm4.crypt_cbc(self.iv, raw)
            unpadded = _pkcs7_unpad(decrypted_bytes)
            return unpadded.decode('utf-8')
        except Exception:
            return value

    def to_python(self, value):
        if value is None:
            return value
        if isinstance(value, bytes):
            try:
                value = value.decode('utf-8')
            except Exception:
                return value
        if isinstance(value, str):
            try:
                raw = base64.b64decode(value)
                crypt_sm4 = CryptSM4()
                crypt_sm4.set_key(self.key, SM4_DECRYPT)
                decrypted_bytes = crypt_sm4.crypt_cbc(self.iv, raw)
                unpadded = _pkcs7_unpad(decrypted_bytes)
                return unpadded.decode('utf-8')
            except Exception:
                return value
        return value

    def get_prep_value(self, value):
        if value is None:
            return value
        if not isinstance(value, str):
            value = str(value)
        data = value.encode('utf-8')
        data = _pkcs7_pad(data)
        crypt_sm4 = CryptSM4()
        crypt_sm4.set_key(self.key, SM4_ENCRYPT)
        encrypted_bytes = crypt_sm4.crypt_cbc(self.iv, data)
        return base64.b64encode(encrypted_bytes).decode('utf-8')