# -*- coding: utf-8 -*-
import base64
import struct

delta = 0x9e3779b9

TEAHexKey = "33576f37387276645054657a6831686e454664424e5652456f47794835795942"

TEAKey = bytes.fromhex(TEAHexKey).decode('utf-8')


class NTea:

    @staticmethod
    def to_tea_key(key_str: str):
        key_bytes = key_str.encode('utf-8')
        key = [0] * 4
        for i in range(4):
            key[i] = struct.unpack('>I', key_bytes[i * 4:(i + 1) * 4])[0]
        return key

    @staticmethod
    def encrypt(v: list[int], tea_key: list[int]):
        """TEA 加密"""
        sum = 0
        n = 32  # 32轮加密
        for _ in range(n):
            sum = (sum + delta) & 0xffffffff
            v[0] = (v[0] + (((v[1] << 4) + tea_key[0]) ^ (v[1] + sum) ^
                            ((v[1] >> 5) + tea_key[1]))) & 0xffffffff
            v[1] = (v[1] + (((v[0] << 4) + tea_key[2]) ^ (v[0] + sum) ^
                            ((v[0] >> 5) + tea_key[3]))) & 0xffffffff
        return v

    @staticmethod
    def decrypt(v: list[int], tea_key: list[int]):
        """TEA 解密"""
        n = 32  # 32轮解密
        sum = (delta * n) & 0xffffffff
        for _ in range(n):
            v[1] = (v[1] - (((v[0] << 4) + tea_key[2]) ^ (v[0] + sum) ^
                            ((v[0] >> 5) + tea_key[3]))) & 0xffffffff
            v[0] = (v[0] - (((v[1] << 4) + tea_key[0]) ^ (v[1] + sum) ^
                            ((v[1] >> 5) + tea_key[1]))) & 0xffffffff
            sum = (sum - delta) & 0xffffffff
        return v

    @staticmethod
    def pkcs7_padding(plaintext: bytes, block_size=8):
        """PKCS#7 填充"""
        padding = block_size - (len(plaintext) % block_size)
        return plaintext + bytes([padding] * padding)

    @staticmethod
    def pkcs7_unpadding(ciphertext: bytes):
        """PKCS#7 去除填充"""
        padding = ciphertext[-1]
        return ciphertext[:-padding]

    @staticmethod
    def bytes_to_uint32(bytes_seq: bytes) -> list[int]:
        return list(struct.unpack('>II', bytes_seq))

    @staticmethod
    def uint32_to_bytes(uint32_seq: list[int]):
        return struct.pack('>II', *uint32_seq)

    @staticmethod
    def encode(plaintext: str, key: str = TEAKey):
        padded = NTea.pkcs7_padding(plaintext.encode("utf-8"))
        tea_key = NTea.to_tea_key(key)

        encrypted_blocks = []

        for i in range(0, len(padded), 8):
            block = NTea.bytes_to_uint32(padded[i:i + 8])
            encrypted_block = NTea.encrypt(block, tea_key)
            encrypted_blocks.append(NTea.uint32_to_bytes(encrypted_block))

        encrypted = b''.join(encrypted_blocks)
        return base64.b64encode(encrypted).decode('utf-8')

    @staticmethod
    def decode(ciphertext: str, key: str = TEAKey):
        cipherbytes = base64.b64decode(ciphertext)
        tea_key = NTea.to_tea_key(key)
        decryptedBlocks = []

        for i in range(0, len(cipherbytes), 8):
            block = NTea.bytes_to_uint32(cipherbytes[i:i + 8])
            decryptedBlock = NTea.decrypt(block, tea_key)
            decryptedBlocks.append(NTea.uint32_to_bytes(decryptedBlock))

        decrypted = b''.join(decryptedBlocks)
        return NTea.pkcs7_unpadding(decrypted).decode("utf-8")
