import numpy as np
from PIL import Image, PngImagePlugin
import time
import hashlib
from typing import Optional, Tuple


def _pad_to_even_square(img: Image.Image) -> Tuple[Image.Image, int, int, int, int]:
    """
    将图像填充为正方形，且边长为偶数。
    返回：(填充后的图像, 原始宽度, 原始高度, 左填充量, 上填充量)
    """
    width, height = img.size
    side = max(width, height)
    if side % 2 != 0:
        side += 1  # 确保为偶数
    delta_w = side - width
    delta_h = side - height
    left = delta_w // 2
    top = delta_h // 2
    # 右和下填充可能多一个像素（当差值不是偶数时）
    right = delta_w - left
    bottom = delta_h - top
    padded_img = Image.new("RGB", (side, side), (128, 128, 128))
    padded_img.paste(img, (left, top))
    return padded_img, width, height, left, top


def _bakers_forward(pixels: np.ndarray, n: int) -> np.ndarray:
    """
    离散 Baker 映射正向变换（混淆）。
    pixels: (n, n, 3) 的 numpy 数组，n 为偶数。
    """
    new_pixels = np.zeros_like(pixels)
    half = n // 2
    for y in range(n):
        for x in range(n):
            if x < half:
                new_x = 2 * x
                new_y = y
            else:
                new_x = 2 * x - n + 1
                new_y = y + half
            # 理论上 new_x, new_y 都在 [0, n-1] 内，取模仅为边界安全
            new_pixels[new_y % n, new_x % n] = pixels[y, x]
    return new_pixels


def _bakers_inverse(pixels: np.ndarray, n: int) -> np.ndarray:
    """
    离散 Baker 映射逆向变换（恢复）。
    pixels: (n, n, 3) 的 numpy 数组，n 为偶数。
    """
    new_pixels = np.zeros_like(pixels)
    half = n // 2
    for y in range(n):
        for x in range(n):
            if x % 2 == 0:
                orig_x = x // 2
                orig_y = y
            else:
                orig_x = (x + n - 1) // 2
                orig_y = y - half
            # 对于正确生成的混淆图像，orig_x, orig_y 均在 [0, n-1] 内
            new_pixels[orig_y % n, orig_x % n] = pixels[y, x]
    return new_pixels


def _derive_key(password: Optional[str], seed: Optional[int]) -> Optional[bytes]:
    """
    根据密码和种子派生密钥（SHA256）。
    如果两者均为 None，则返回 None。
    """
    if password is None and seed is None:
        return None
    material = ''
    if password:
        material += password
    if seed is not None:
        material += str(seed)
    return hashlib.sha256(material.encode()).digest()


def _apply_xor(pixels: np.ndarray, key: bytes) -> np.ndarray:
    """
    对像素数组应用逐字节异或（原地修改并返回）。
    pixels: 形状为 (H, W, C) 的 uint8 数组。
    key: 字节密钥。
    """
    key_array = np.frombuffer(key, dtype=np.uint8)
    # 将密钥扩展为像素数组大小
    repeats = (pixels.size + key_array.size - 1) // key_array.size
    key_expanded = np.tile(key_array, repeats)[:pixels.size]
    key_expanded = key_expanded.reshape(pixels.shape)
    # 异或操作（numpy 自动按元素进行）
    return np.bitwise_xor(pixels, key_expanded)


def obfuscate(
    image_path: str,
    output_path: str,
    iterations: int = 20,
    seed: Optional[int] = None,
    password: Optional[str] = None,
    use_timestamp: bool = True
) -> Optional[int]:
    """
    对图像进行 Baker 混淆（可选加密层），并保存为 PNG（含元数据）。
    - image_path: 输入图像路径
    - output_path: 输出混淆图像路径（PNG 格式）
    - iterations: 混淆迭代次数
    - seed: 种子，与 password 一起用于派生加密密钥（可选）
    - password: 密码，用于派生加密密钥（可选）
    - use_timestamp: 是否将当前时间戳嵌入元数据
    返回：嵌入的时间戳（若 use_timestamp=True），否则 None
    """
    img = Image.open(image_path).convert("RGB")
    padded_img, orig_width, orig_height, left_pad, top_pad = _pad_to_even_square(img)
    n = padded_img.size[0]  # 此时 n 为偶数

    pixels = np.array(padded_img)

    # 若提供了密码或种子，则对像素进行异或加密
    key = _derive_key(password, seed)
    if key is not None:
        pixels = _apply_xor(pixels, key)
        encrypted_flag = "1"
    else:
        encrypted_flag = "0"

    timestamp = int(time.time()) if use_timestamp else None

    for _ in range(iterations):
        pixels = _bakers_forward(pixels, n)

    scrambled_img = Image.fromarray(pixels)

    # 保存元数据
    info = PngImagePlugin.PngInfo()
    if timestamp is not None:
        info.add_text("bakers_timestamp", str(timestamp))
    info.add_text("bakers_iterations", str(iterations))
    info.add_text("orig_width", str(orig_width))
    info.add_text("orig_height", str(orig_height))
    info.add_text("padded_size", str(n))
    info.add_text("left_pad", str(left_pad))
    info.add_text("top_pad", str(top_pad))
    info.add_text("bakers_encrypted", encrypted_flag)  # 标记是否加密

    scrambled_img.save(output_path, "PNG", pnginfo=info)
    print(f"Baker 混淆完成，已保存到：{output_path}")
    print(f"迭代次数: {iterations}，填充后边长: {n}")
    if key is not None:
        print("加密层已应用（使用 password/seed）")
    if timestamp is not None:
        print(f"时间戳已嵌入: {timestamp}")

    return timestamp


def deobfuscate(
    obf_image_path: str,
    output_path: str,
    iterations: Optional[int] = None,
    seed: Optional[int] = None,
    timestamp: Optional[int] = None,
    password: Optional[str] = None,
    orig_width: Optional[int] = None,
    orig_height: Optional[int] = None,
    padded_size: Optional[int] = None,
    left_pad: Optional[int] = None,
    top_pad: Optional[int] = None,
    encrypted_flag: Optional[str] = None
) -> None:
    """
    恢复 Baker 混淆的图像（若加密则需要提供相同的 password/seed）。
    - obf_image_path: 混淆图像路径（PNG，应包含元数据）
    - output_path: 输出恢复图像路径
    - iterations: 若未提供，则从元数据读取
    - seed: 种子，用于派生解密密钥（若图像加密）
    - timestamp: 未使用（仅用于接口兼容）
    - password: 密码，用于派生解密密钥（若图像加密）
    - orig_width: 原始图像宽度（pipeline 中间层传入，优先于元数据）
    - orig_height: 原始图像高度（pipeline 中间层传入，优先于元数据）
    - padded_size: 填充后边长（pipeline 中间层传入，优先于元数据）
    - left_pad: 左侧填充量（pipeline 中间层传入，优先于元数据）
    - top_pad: 顶部填充量（pipeline 中间层传入，优先于元数据）
    - encrypted_flag: 加密标记（pipeline 中间层传入，优先于元数据）
    """
    img = Image.open(obf_image_path).convert("RGB")
    pixels = np.array(img)
    n = pixels.shape[0]

    # 从元数据读取参数，若调用方已传入则优先使用（pipeline 中间层无元数据时的兼容）
    iters_str = img.info.get("bakers_iterations")
    if iterations is not None:
        final_iterations = iterations
    elif iters_str is not None:
        final_iterations = int(iters_str)
    else:
        raise ValueError("无法从图像元数据中读取迭代次数，请手动传入 iterations 参数")

    # 优先使用传入的参数，否则从元数据读取
    final_orig_width = orig_width if orig_width is not None else int(img.info.get("orig_width", n))
    final_orig_height = orig_height if orig_height is not None else int(img.info.get("orig_height", n))
    final_padded_size = padded_size if padded_size is not None else int(img.info.get("padded_size", n))
    final_left_pad = left_pad if left_pad is not None else int(img.info.get("left_pad", 0))
    final_top_pad = top_pad if top_pad is not None else int(img.info.get("top_pad", 0))
    final_encrypted_flag = encrypted_flag if encrypted_flag is not None else img.info.get("bakers_encrypted", "0")

    # 边长校验
    if n != final_padded_size:
        print(f"警告: 当前图像边长 {n} 与元数据记录的填充边长 {final_padded_size} 不符，将使用当前边长 {n} 进行恢复")

    ts_from_meta = img.info.get("bakers_timestamp")
    print(f"恢复参数：iterations={final_iterations}, 元数据时间戳={ts_from_meta}, 加密标记={final_encrypted_flag}")

    # 逆向 Baker 迭代
    for _ in range(final_iterations):
        pixels = _bakers_inverse(pixels, n)

    # 如果图像加密了，则需要解密
    if final_encrypted_flag == "1":
        key = _derive_key(password, seed)
        if key is None:
            raise ValueError("图像已加密，但未提供 password 或 seed，无法解密")
        pixels = _apply_xor(pixels, key)
        print("解密层已应用")

    # 根据保存的填充偏移量裁剪原始图像区域
    restored_pixels = pixels[final_top_pad:final_top_pad + final_orig_height, final_left_pad:final_left_pad + final_orig_width, :]
    restored_img = Image.fromarray(restored_pixels)
    restored_img.save(output_path, "PNG")
    print(f"Baker 恢复完成，已保存到：{output_path}")
    print(f"最终尺寸: {final_orig_width} x {final_orig_height}")


# 可选：添加一个简单的测试函数，验证正向和逆向变换是否互逆（不包含加密层）
def test_bakers_map(n: int = 10):
    """测试 n 为偶数时的互逆性"""
    if n % 2 != 0:
        n += 1
    original = np.random.randint(0, 256, (n, n, 3), dtype=np.uint8)
    forward = _bakers_forward(original, n)
    inverse = _bakers_inverse(forward, n)
    if np.array_equal(original, inverse):
        print("测试通过：正向和逆向 Baker 变换互逆")
    else:
        print("测试失败：正向和逆向 Baker 变换不互逆")
        diff = np.sum(original != inverse)
        print(f"差异像素数：{diff}")


if __name__ == "__main__":
    # 测试 Baker 映射本身（可选）
    # test_bakers_map(128)

    # 使用示例（混淆 + 密码）

    pass