# -*- coding: utf-8 -*-
"""
文件名加密模块

提供文件名加密功能，严格遵循 clouddrive-decrypt 标准
"""

from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad, unpad
from Crypto.Hash import HMAC, SHA1
from Crypto.Protocol.KDF import PBKDF2


class FilenameEncryptor:
    """文件名加密器（遵循 clouddrive-decrypt 标准）"""
    
    def __init__(self, key: bytes, salt: bytes):
        """
        初始化加密器
        
        Args:
            key: 加密密钥（密码的字节形式）
            salt: 盐值（8 字节）
        """
        if len(salt) != 8:
            raise ValueError(f"Salt 长度必须为 8 字节，当前：{len(salt)}")
        self.key = key
        self.salt = salt
    
    def encrypt(self, filename: str) -> str:
        """
        加密文件名
        
        过程：
        1. 将文件名编码为 UTF-8
        2. 使用 PBKDF2-HMAC-SHA1 派生密钥（1000 次迭代，8 字节 salt）
        3. 使用 AES-256-CBC 加密
        4. 使用自定义 UTF-8 编码（中文字符）
        5. 添加 .cdcrypto 扩展名
        
        Args:
            filename: 原始文件名
        
        Returns:
            str: 加密后的文件名
        """
        utf8_bytes = filename.encode('utf-8')
        
        # 使用 PBKDF2-HMAC-SHA1 派生密钥（1000 次迭代）
        kdf = PBKDF2(
            self.key,
            self.salt,
            count=1000,
            dkLen=32 + 16,  # 32 字节密钥 + 16 字节 IV
            hmac_hash_module=SHA1
        )
        derived_key = kdf[:32]
        derived_iv = kdf[32:]
        
        # AES-256-CBC 加密
        cipher = AES.new(derived_key, AES.MODE_CBC, derived_iv)
        padded_data = pad(utf8_bytes, 16)
        encrypted = cipher.encrypt(padded_data)
        
        # 拼接 salt + 加密数据
        encrypted_with_salt = self.salt + encrypted
        
        # 使用自定义 UTF-8 编码（中文字符）
        encoded_str = self._base_utf8_encode(encrypted_with_salt)
        
        # 添加 .cdcrypto 扩展名
        encrypted_name = f"{encoded_str}.cdcrypto"
        
        return encrypted_name
    
    def decrypt(self, encrypted_name: str) -> str:
        """
        解密文件名
        
        过程：
        1. 移除 .cdcrypto 扩展名
        2. 自定义 UTF-8 解码（中文字符）
        3. 提取 salt（前 8 字节）
        4. 使用 PBKDF2-HMAC-SHA1 派生密钥（1000 次迭代）
        5. AES-256-CBC 解密
        6. 移除填充
        7. UTF-8 解码
        
        Args:
            encrypted_name: 加密后的文件名
        
        Returns:
            str: 原始文件名
        """
        if not encrypted_name.endswith('.cdcrypto'):
            raise ValueError("文件名必须以 .cdcrypto 结尾")
        
        encoded_part = encrypted_name[:-9]
        
        # 自定义 UTF-8 解码（中文字符）
        encrypted_with_salt = self._base_utf8_decode(encoded_part)
        
        if len(encrypted_with_salt) <= 8:
            raise ValueError("加密数据长度无效")
        
        # 提取 salt 和加密数据
        salt = encrypted_with_salt[:8]
        encrypted_data = encrypted_with_salt[8:]
        
        # 使用 PBKDF2-HMAC-SHA1 派生密钥（1000 次迭代）
        kdf = PBKDF2(
            self.key,
            salt,
            count=1000,
            dkLen=32 + 16,  # 32 字节密钥 + 16 字节 IV
            hmac_hash_module=SHA1
        )
        derived_key = kdf[:32]
        derived_iv = kdf[32:]
        
        # AES-256-CBC 解密
        cipher = AES.new(derived_key, AES.MODE_CBC, derived_iv)
        decrypted = cipher.decrypt(encrypted_data)
        unpadded = unpad(decrypted, 16)
        
        # UTF-8 解码
        filename = unpadded.decode('utf-8')
        
        return filename
    
    def _base_utf8_encode(self, data: bytes) -> str:
        """
        自定义 UTF-8 编码（使用中文字符）
        
        参考 base_utf8_decode.rs 的 decode 函数的反向过程。
        将二进制数据编码为中文字符串，每个字符代表 14 位。
        
        Args:
            data: 原始二进制数据
        
        Returns:
            str: 编码后的中文字符串
        """
        output = []
        bit_buffer = 0
        bit_buffer_len = 0
        
        # 起始 Unicode 码点（CJK Unified Ideographs）
        start = 0x4E00
        
        # 前 14 个字符是填充字符
        MAX_PADDING = 14
        
        data_len = len(data)
        pos = 0
        
        while pos < data_len:
            # 从数据中读取字节到缓冲区
            byte = data[pos]
            bit_buffer = (bit_buffer << 8) | byte
            bit_buffer_len += 8
            pos += 1
            
            # 当缓冲区有足够位数时，输出字符
            while bit_buffer_len >= 14:
                bit_buffer_len -= 14
                index = ((bit_buffer >> bit_buffer_len) & 0x3FFF)
                char = chr(start + index)
                output.append(char)
                bit_buffer &= (1 << bit_buffer_len) - 1
        
        # 处理剩余位（添加填充）
        if bit_buffer_len > 0:
            # 左对齐剩余位
            index = bit_buffer << (14 - bit_buffer_len)
            # 填充索引 = 14 - 剩余位数
            padding_index = 14 - bit_buffer_len
            # 将剩余数据编码到最后一个字符
            last_index = index | padding_index
            char = chr(start + last_index)
            output.append(char)
        
        return ''.join(output)
    
    def _base_utf8_decode(self, s: str) -> bytes:
        """
        自定义 UTF-8 解码（使用中文字符）
        
        参考 base_utf8_decode.rs 的 decode 函数。
        将中文字符串解码为二进制数据。
        
        Args:
            s: 编码后的中文字符串
        
        Returns:
            bytes: 解码后的二进制数据
        """
        output = bytearray()
        bit_buffer = 0
        bit_buffer_len = 0
        
        # 起始 Unicode 码点（CJK Unified Ideographs）
        start = 0x4E00
        
        # 前 14 个字符是填充字符
        MAX_PADDING = 14
        
        chars = list(s)
        
        # 检查最后一个字符是否是填充字符
        padding_bits = 0
        has_padding = False
        data_len = len(chars)
        
        if chars:
            last_char = chars[-1]
            last_code = ord(last_char) - start
            if 0 <= last_code < MAX_PADDING:
                # 最后一个字符是填充字符
                padding_bits = last_code + 1
                has_padding = True
                data_len -= 1
        
        # 解码数据字符
        for i, char in enumerate(chars[:data_len]):
            code = ord(char) - start
            if code < 0:
                continue
            
            if has_padding and i == data_len - 1:
                # 最后一个字符有填充
                valid_bits = 14 - padding_bits
                data_bits = code >> padding_bits
                bit_buffer = (bit_buffer << valid_bits) | data_bits
                bit_buffer_len += valid_bits
            else:
                bit_buffer = (bit_buffer << 14) | code
                bit_buffer_len += 14
            
            # 从缓冲区提取字节
            while bit_buffer_len >= 8:
                bit_buffer_len -= 8
                byte = ((bit_buffer >> bit_buffer_len) & 0xFF)
                output.append(byte)
                bit_buffer &= (1 << bit_buffer_len) - 1
        
        return bytes(output)


def create_encryptor(password: str) -> FilenameEncryptor:
    """
    创建加密器实例（遵循 clouddrive-decrypt 标准）
    
    使用随机生成的 8 字节 salt。
    
    Args:
        password: 加密密码
    
    Returns:
        FilenameEncryptor: 加密器实例
    """
    salt = get_random_bytes(8)
    return FilenameEncryptor(password.encode('utf-8'), salt)
