"""
==============================================================================
ED25519 消息加密模块 (crypto.py)
==============================================================================

本模块实现基于 AE 账户的消息加密/解密功能。
利用 ED25519 -> X25519 密钥转换，实现 NaCl 公钥加密。

==============================================================================
技术原理
==============================================================================

1. 密钥转换:
   - AE 账户使用 ED25519 密钥对 (用于签名)
   - NaCl Box 加密使用 X25519 密钥对 (用于加密)
   - 两者基于相同的 Curve25519 曲线，可以数学转换

2. 加密方案:
   
   方案 A: SealedBox (匿名加密)
   ┌─────────────────────────────────────────────────────────────┐
   │  发送方: 只需 [接收方公钥]                                  │
   │  - 生成临时密钥对                                           │
   │  - 用临时私钥 + 接收方公钥 加密                             │
   │  - 临时私钥随即销毁                                         │
   │                                                              │
   │  接收方: 用 [自己的私钥] 解密                               │
   │  - 无法知道发送者身份                                       │
   │  - 适合匿名消息、公开投递                                   │
   └─────────────────────────────────────────────────────────────┘
   
   方案 B: Box (双向认证加密)
   ┌─────────────────────────────────────────────────────────────┐
   │  发送方: 用 [自己的私钥 + 接收方公钥] 加密                  │
   │  接收方: 用 [自己的私钥 + 发送方公钥] 解密                  │
   │                                                              │
   │  - 双方都能验证对方身份                                     │
   │  - 适合双向通信、身份认证场景                               │
   └─────────────────────────────────────────────────────────────┘

3. 密文格式:
   - SealedBox: 明文 + 48 字节 (32 字节临时公钥 + 16 字节 MAC)
   - Box: 明文 + 40 字节 (24 字节 nonce + 16 字节 MAC)

==============================================================================
"""

import base64

from nacl.public import PrivateKey as X25519PrivateKey, PublicKey as X25519PublicKey
from nacl.public import Box, SealedBox
from nacl.signing import SigningKey, VerifyKey
from nacl.encoding import RawEncoder, HexEncoder

from aeknow import signing, hashing


# =============================================================================
# 密钥转换函数
# =============================================================================

def ed25519_to_x25519_private(signing_key: SigningKey) -> X25519PrivateKey:
    """
    将 ED25519 签名私钥转换为 X25519 加密私钥
    
    :param signing_key: NaCl SigningKey (ED25519)
    :return: NaCl X25519 PrivateKey
    
    原理:
        ED25519 私钥 (32 字节种子) -> SHA512 -> 前 32 字节 -> X25519 私钥
        PyNaCl 内置了这个转换函数
    """
    return signing_key.to_curve25519_private_key()


def ed25519_to_x25519_public(verify_key: VerifyKey) -> X25519PublicKey:
    """
    将 ED25519 验证公钥转换为 X25519 加密公钥
    
    :param verify_key: NaCl VerifyKey (ED25519)
    :return: NaCl X25519 PublicKey
    
    原理:
        ED25519 公钥点 (在 Ed25519 曲线上) -> 坐标转换 -> X25519 公钥点
        这是一个双向映射 (birational map)
    """
    return verify_key.to_curve25519_public_key()


def account_to_x25519_keypair(account: signing.Account) -> tuple:
    """
    将 AE Account 的 ED25519 密钥对转换为 X25519 密钥对
    
    :param account: AE Account 对象
    :return: (x25519_private_key, x25519_public_key)
    """
    x25519_private = ed25519_to_x25519_private(account.signing_key)
    x25519_public = ed25519_to_x25519_public(account.verifying_key)
    return x25519_private, x25519_public


def address_to_x25519_public(address: str) -> X25519PublicKey:
    """
    从 AE 地址 (ak_xxx) 提取并转换为 X25519 公钥
    
    :param address: AE 账户地址 (ak_xxx 格式)
    :return: X25519 PublicKey
    
    流程:
        1. ak_xxx -> base58check 解码 -> 32 字节 ED25519 公钥
        2. ED25519 公钥 -> X25519 公钥
    """
    # 解码 ak_xxx 地址得到原始公钥字节
    ed25519_pubkey_bytes = hashing.decode(address)
    
    # 创建 VerifyKey 对象
    verify_key = VerifyKey(ed25519_pubkey_bytes, encoder=RawEncoder)
    
    # 转换为 X25519 公钥
    return ed25519_to_x25519_public(verify_key)


# =============================================================================
# SealedBox 匿名加密 (只需接收方公钥)
# =============================================================================

def sealed_encrypt(message: str, recipient_address: str) -> bytes:
    """
    使用 SealedBox 匿名加密消息
    
    :param message: 明文消息 (UTF-8 字符串)
    :param recipient_address: 接收方 AE 地址 (ak_xxx)
    :return: 密文 (bytes)
    
    特点:
        - 发送方匿名，接收方无法知道消息来源
        - 只需要接收方的公钥
        - 密文比明文长 48 字节
    """
    # 获取接收方的 X25519 公钥
    recipient_x25519_pub = address_to_x25519_public(recipient_address)
    
    # 创建 SealedBox (只需接收方公钥)
    sealed_box = SealedBox(recipient_x25519_pub)
    
    # 加密消息
    msg_bytes = message.encode('utf-8')
    ciphertext = sealed_box.encrypt(msg_bytes)
    
    return ciphertext


def sealed_decrypt(ciphertext: bytes, account: signing.Account) -> str:
    """
    使用 SealedBox 解密匿名消息
    
    :param ciphertext: 密文 (bytes)
    :param account: 接收方 AE Account 对象
    :return: 解密后的明文消息
    """
    # 转换接收方私钥为 X25519
    x25519_private, _ = account_to_x25519_keypair(account)
    
    # 创建 SealedBox (用私钥解密)
    unseal_box = SealedBox(x25519_private)
    
    # 解密
    plaintext = unseal_box.decrypt(ciphertext)
    
    return plaintext.decode('utf-8')


# =============================================================================
# Box 双向认证加密 (需要双方密钥)
# =============================================================================

def box_encrypt(message: str, sender_account: signing.Account, recipient_address: str) -> bytes:
    """
    使用 Box 双向认证加密消息
    
    :param message: 明文消息 (UTF-8 字符串)
    :param sender_account: 发送方 AE Account 对象
    :param recipient_address: 接收方 AE 地址 (ak_xxx)
    :return: 密文 (bytes)，包含 nonce
    
    特点:
        - 接收方可以验证消息确实来自发送方
        - 需要发送方私钥 + 接收方公钥
        - 密文比明文长 40 字节 (24 nonce + 16 MAC)
    """
    # 转换发送方私钥
    sender_x25519_private, _ = account_to_x25519_keypair(sender_account)
    
    # 获取接收方 X25519 公钥
    recipient_x25519_pub = address_to_x25519_public(recipient_address)
    
    # 创建 Box (发送方私钥 + 接收方公钥)
    box = Box(sender_x25519_private, recipient_x25519_pub)
    
    # 加密消息 (自动生成随机 nonce)
    msg_bytes = message.encode('utf-8')
    encrypted = box.encrypt(msg_bytes)  # 包含 nonce + ciphertext
    
    return bytes(encrypted)


def box_decrypt(ciphertext: bytes, recipient_account: signing.Account, sender_address: str) -> str:
    """
    使用 Box 双向认证解密消息
    
    :param ciphertext: 密文 (bytes)，包含 nonce
    :param recipient_account: 接收方 AE Account 对象
    :param sender_address: 发送方 AE 地址 (ak_xxx)
    :return: 解密后的明文消息
    
    注意:
        - 需要发送方地址来验证消息来源
        - 如果发送方地址不匹配，解密会失败
    """
    # 转换接收方私钥
    recipient_x25519_private, _ = account_to_x25519_keypair(recipient_account)
    
    # 获取发送方 X25519 公钥
    sender_x25519_pub = address_to_x25519_public(sender_address)
    
    # 创建 Box (接收方私钥 + 发送方公钥)
    box = Box(recipient_x25519_private, sender_x25519_pub)
    
    # 解密消息
    plaintext = box.decrypt(ciphertext)
    
    return plaintext.decode('utf-8')


# =============================================================================
# 便捷封装类
# =============================================================================

class MessageCrypto:
    """
    消息加密工具类
    
    封装了账户加载和加密/解密操作
    """
    
    def __init__(self, account: signing.Account = None):
        self.account = account
        self.x25519_private = None
        self.x25519_public = None
        if account:
            self._init_x25519_keys()
    
    def _init_x25519_keys(self):
        """初始化 X25519 密钥对"""
        if self.account:
            self.x25519_private, self.x25519_public = account_to_x25519_keypair(self.account)
    
    @classmethod
    def from_keystore(cls, path: str, password: str) -> 'MessageCrypto':
        """从 Keystore 加载账户"""
        account = signing.Account.from_keystore(path, password)
        return cls(account)
    
    @classmethod
    def from_mnemonic(cls, mnemonic: str) -> 'MessageCrypto':
        """从助记词加载账户"""
        from aeknow.hdwallet import HDWallet
        wallet = HDWallet(mnemonic)
        return cls(wallet.get_master_account())
    
    @classmethod
    def from_account(cls, account: signing.Account) -> 'MessageCrypto':
        """从已有 Account 对象加载"""
        return cls(account)
    
    @property
    def address(self) -> str:
        """获取当前账户地址"""
        return self.account.get_address() if self.account else None
    
    @property
    def x25519_public_hex(self) -> str:
        """获取 X25519 公钥的十六进制表示"""
        if self.x25519_public:
            return self.x25519_public.encode(encoder=HexEncoder).decode()
        return None
    
    # -------------------------------------------------------------------------
    # SealedBox 匿名加密
    # -------------------------------------------------------------------------
    
    def sealed_encrypt(self, message: str, recipient_address: str) -> str:
        """
        匿名加密消息 (发送给他人)
        
        :param message: 明文消息
        :param recipient_address: 接收方地址 ak_xxx
        :return: Base64 编码的密文
        """
        ciphertext = sealed_encrypt(message, recipient_address)
        return base64.b64encode(ciphertext).decode('ascii')
    
    def sealed_decrypt(self, ciphertext_b64: str) -> str:
        """
        解密匿名消息 (发给自己的)
        
        :param ciphertext_b64: Base64 编码的密文
        :return: 解密后的明文
        """
        if not self.account:
            raise ValueError("未加载账户")
        ciphertext = base64.b64decode(ciphertext_b64)
        return sealed_decrypt(ciphertext, self.account)
    
    # -------------------------------------------------------------------------
    # Box 双向认证加密
    # -------------------------------------------------------------------------
    
    def box_encrypt(self, message: str, recipient_address: str) -> str:
        """
        双向认证加密消息 (对方可验证来源)
        
        :param message: 明文消息
        :param recipient_address: 接收方地址 ak_xxx
        :return: Base64 编码的密文
        """
        if not self.account:
            raise ValueError("未加载账户")
        ciphertext = box_encrypt(message, self.account, recipient_address)
        return base64.b64encode(ciphertext).decode('ascii')
    
    def box_decrypt(self, ciphertext_b64: str, sender_address: str) -> str:
        """
        双向认证解密消息 (验证发送方身份)
        
        :param ciphertext_b64: Base64 编码的密文
        :param sender_address: 发送方地址 ak_xxx
        :return: 解密后的明文
        """
        if not self.account:
            raise ValueError("未加载账户")
        ciphertext = base64.b64decode(ciphertext_b64)
        return box_decrypt(ciphertext, self.account, sender_address)


def get_key_info(account: signing.Account) -> dict:
    """
    获取账户密钥信息
    
    :param account: AE Account 对象
    :return: 包含地址、ED25519 和 X25519 公钥的字典
    """
    x25519_priv, x25519_pub = account_to_x25519_keypair(account)
    
    return {
        'address': account.get_address(),
        'ed25519_public': account.verifying_key.encode(encoder=HexEncoder).decode(),
        'x25519_public': x25519_pub.encode(encoder=HexEncoder).decode()
    }


def print_key_info(account: signing.Account):
    """打印账户的密钥信息"""
    x25519_priv, x25519_pub = account_to_x25519_keypair(account)
    
    print("\n" + "="*70)
    print("账户密钥信息")
    print("="*70)
    print(f"AE 地址 (ak_xxx):        {account.get_address()}")
    print(f"ED25519 公钥 (hex):      {account.verifying_key.encode(encoder=HexEncoder).decode()}")
    print(f"X25519 公钥 (hex):       {x25519_pub.encode(encoder=HexEncoder).decode()}")
    print("="*70 + "\n")
