#!/usr/bin/env python3
"""
腾讯混元3D签名算法实现
从 webpack 模块逆向提取
"""

import hmac
import hashlib
import time
import random
import string

# 密钥派生相关常量（从 JS 提取）
C = bytes([122, 59, 92, 165, 30, 79, 166, 139, 142, 129, 139, 89, 219, 131, 101, 204])
D = bytes([122, 59, 92, 45, 30, 79, 106, 139, 156, 13, 46, 63, 74, 91, 108, 125])
U = [3, 5, 2, 7, 1, 4, 6, 2, 5, 3, 1, 4, 2, 6, 3, 5]
M = [14, 11, 13, 9, 15, 10, 12, 8, 6, 3, 5, 1, 7, 2, 4, 0]


def derive_key(c: bytes) -> str:
    """密钥派生函数 - 从硬编码常量派生签名密钥"""
    if len(c) != 16:
        raise ValueError("输入必须是一个16字节的数组")

    # 步骤1: XOR with D
    t = bytearray(16)
    for i in range(16):
        t[i] = c[i] ^ D[i]

    # 步骤2: 循环左移
    o = bytearray(16)
    for i in range(16):
        n = U[i]
        # Python 的移位和 JS 不同，需要模拟 8 位无符号整数
        val = t[i]
        o[i] = (val << n | val >> (8 - n)) & 0xFF

    # 步骤3: 置换
    n = bytearray(16)
    for i in range(16):
        n[i] = o[M[i]]

    # 步骤4: 找到第一个0字节，截断
    try:
        r = n.index(0)
    except ValueError:
        r = 16

    return n[:r].decode('utf-8')


def generate_nonce(length: int = 16) -> str:
    """生成随机 nonce"""
    chars = string.ascii_letters + string.digits  # 62字符
    return ''.join(random.choice(chars) for _ in range(length))


def get_timestamp() -> int:
    """获取当前时间戳（秒）"""
    return int(time.time())


def sort_params(params: dict) -> list:
    """排序参数，过滤空值"""
    import json
    items = []
    for k, v in params.items():
        if v is not None and v != "":
            # 对列表和字典使用 JSON 格式
            if isinstance(v, (list, dict, bool)):
                items.append((k, json.dumps(v, separators=(',', ':'), ensure_ascii=False)))
            else:
                items.append((k, str(v)))
    return sorted(items, key=lambda x: x[0])


def join_params(items: list) -> str:
    """拼接参数为查询字符串"""
    return '&'.join(f"{k}={v}" for k, v in items)


def sign(params: dict,
         nonce_length: int = 16,
         timestamp_field: str = "timestamp",
         nonce_field: str = "nonce",
         sign_field: str = "sign") -> dict:
    """
    主签名函数

    Args:
        params: 请求参数
        nonce_length: nonce 长度
        timestamp_field: 时间戳字段名
        nonce_field: nonce 字段名
        sign_field: 签名字段名

    Returns:
        包含签名的新参数字典
    """
    result = dict(params)
    result[timestamp_field] = get_timestamp()
    result[nonce_field] = generate_nonce(nonce_length)

    # 排序并拼接
    sorted_items = sort_params(result)
    param_str = join_params(sorted_items)

    # 派生密钥
    key = derive_key(C)

    # HMAC-SHA256 签名，然后转为 Hex
    signature = hmac.new(
        key.encode('utf-8'),
        param_str.encode('utf-8'),
        hashlib.sha256
    ).hexdigest()

    result[sign_field] = signature
    return result


def sign_with_custom_nonce(params: dict, timestamp: int, nonce: str,
                           timestamp_field: str = "timestamp",
                           nonce_field: str = "nonce",
                           sign_field: str = "sign") -> dict:
    """使用指定的 timestamp 和 nonce 生成签名（用于验证）"""
    result = dict(params)
    result[timestamp_field] = timestamp
    result[nonce_field] = nonce

    sorted_items = sort_params(result)
    param_str = join_params(sorted_items)

    key = derive_key(C)

    signature = hmac.new(
        key.encode('utf-8'),
        param_str.encode('utf-8'),
        hashlib.sha256
    ).hexdigest()

    result[sign_field] = signature
    return result


if __name__ == "__main__":
    # 测试密钥派生
    key = derive_key(C)
    print(f"派生密钥: {key}")
    print(f"密钥长度: {len(key)}")

    # 测试签名
    test_params = {
        "sceneType": "playGround3D-2.0",
        "count": 1,
        "modelType": "image2ModelV3.1"
    }

    signed = sign(test_params)
    print(f"\n测试签名:")
    for k, v in sorted(signed.items()):
        print(f"  {k}: {v}")