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


def _generate_final_seed(seed: Optional[int], timestamp: Optional[int] = None) -> int:
    parts = []
    if seed is not None:
        parts.append(str(seed))
    if timestamp is not None:
        parts.append(str(timestamp))
    if not parts:
        parts.append("default_seed_1234567890")

    combined = "_".join(parts)
    hash_int = int(hashlib.sha256(combined.encode()).hexdigest(), 16)
    return hash_int % (2 ** 32)


def obfuscate(
    image_path: str,
    output_path: str,
    block_size: int = 16,
    seed: Optional[int] = None,
    use_timestamp: bool = True
) -> Optional[int]:
    """
    块级随机置换（混淆）：将图像分成 block_size×block_size 的块，打乱块顺序，
    并将 timestamp、原始尺寸、填充后尺寸嵌入 PNG 元数据。
    """
    img = Image.open(image_path).convert("RGB")
    orig_width, orig_height = img.size           # Pillow: (width, height)
    pixels = np.array(img)
    height, width, channels = pixels.shape       # numpy: (height, width, channels)

    print(f"原始图像 (Pillow size): {orig_width}x{orig_height}")
    print(f"numpy shape: {height}x{width}x{channels}")

    # 计算 padding
    pad_height = (block_size - height % block_size) % block_size
    pad_width = (block_size - width % block_size) % block_size

    if pad_height > 0 or pad_width > 0:
        pixels = np.pad(pixels, ((0, pad_height), (0, pad_width), (0, 0)), mode='reflect')

    padded_height, padded_width, _ = pixels.shape

    num_rows = padded_height // block_size
    num_cols = padded_width // block_size

    # 分块
    blocks = pixels.reshape(num_rows, block_size, num_cols, block_size, channels)
    blocks = blocks.transpose(0, 2, 1, 3, 4)  # (num_rows, num_cols, bs, bs, ch)
    flat_blocks = blocks.reshape(-1, block_size, block_size, channels)

    # 生成种子
    timestamp = int(time.time()) if use_timestamp else None
    final_seed = _generate_final_seed(seed, timestamp)

    # 打乱
    rng = np.random.RandomState(final_seed)
    indices = np.arange(len(flat_blocks))
    rng.shuffle(indices)
    shuffled_blocks = flat_blocks[indices]

    # 重组
    shuffled_reshaped = shuffled_blocks.reshape(num_rows, num_cols, block_size, block_size, channels)
    shuffled_pixels = shuffled_reshaped.transpose(0, 2, 1, 3, 4).reshape(padded_height, padded_width, channels)

    shuffled_img = Image.fromarray(shuffled_pixels)

    # 保存元数据
    info = PngImagePlugin.PngInfo()
    if timestamp is not None:
        info.add_text("obf_timestamp", str(timestamp))
    info.add_text("orig_width", str(orig_width))
    info.add_text("orig_height", str(orig_height))
    info.add_text("padded_width", str(padded_width))
    info.add_text("padded_height", str(padded_height))

    shuffled_img.save(output_path, "PNG", pnginfo=info)
    print(f"图像混淆成功，已保存到：{output_path}")
    print(f"原始尺寸: {orig_width}x{orig_height} → padded: {padded_width}x{padded_height}")
    if timestamp is not None:
        print(f"timestamp 已嵌入: {timestamp}")

    return timestamp


def deobfuscate(
    obf_image_path: str,
    output_path: str,
    block_size: int = 16,
    seed: Optional[int] = None,
    timestamp: Optional[int] = None
) -> None:
    """
    恢复块级混淆图像：从元数据读取原始尺寸和填充尺寸，确保裁剪正确。
    """
    img = Image.open(obf_image_path).convert("RGB")
    pixels = np.array(img)
    current_height, current_width, channels = pixels.shape

    # 读取元数据
    orig_width = int(img.info.get("orig_width", current_width))
    orig_height = int(img.info.get("orig_height", current_height))
    padded_width = int(img.info.get("padded_width", current_width))
    padded_height = int(img.info.get("padded_height", current_height))

    print(f"从元数据读取：原始 {orig_width}x{orig_height}，padded {padded_width}x{padded_height}")

    # timestamp 处理
    ts_from_meta = None
    ts_str = img.info.get("obf_timestamp")
    if ts_str:
        try:
            ts_from_meta = int(ts_str)
            print(f"从元数据读取 timestamp: {ts_from_meta}")
        except ValueError:
            print("警告：obf_timestamp 格式无效")

    final_timestamp = ts_from_meta if ts_from_meta is not None else timestamp
    if final_timestamp is None:
        print("未找到 timestamp，使用纯 seed 模式")

    final_seed = _generate_final_seed(seed, final_timestamp)

    # 检查当前尺寸是否匹配 padded
    if current_height != padded_height or current_width != padded_width:
        print(f"警告：当前加载尺寸 {current_width}x{current_height} 与元数据 padded {padded_width}x{padded_height} 不符")
        print("可能原因：元数据丢失或图像被重新保存导致尺寸变化")
        # 强制使用当前尺寸继续（可能导致部分内容丢失）
        padded_width = current_width
        padded_height = current_height

    num_rows = padded_height // block_size
    num_cols = padded_width // block_size

    if num_rows * block_size != padded_height or num_cols * block_size != padded_width:
        raise ValueError(
            f"填充后尺寸 {padded_width}x{padded_height} 无法被 block_size={block_size} 整除"
        )

    # 分块
    blocks = pixels.reshape(num_rows, block_size, num_cols, block_size, channels)
    blocks = blocks.transpose(0, 2, 1, 3, 4)
    flat_blocks = blocks.reshape(-1, block_size, block_size, channels)

    # 逆置换
    rng = np.random.RandomState(final_seed)
    indices = np.arange(len(flat_blocks))
    rng.shuffle(indices)
    inverse_indices = np.argsort(indices)
    restored_blocks = flat_blocks[inverse_indices]

    # 重组
    restored_reshaped = restored_blocks.reshape(num_rows, num_cols, block_size, block_size, channels)
    restored_pixels = restored_reshaped.transpose(0, 2, 1, 3, 4).reshape(padded_height, padded_width, channels)

    # 裁剪回原始尺寸（注意顺序：高度在前，宽度在后）
    restored_pixels = restored_pixels[:orig_height, :orig_width, :]

    restored_img = Image.fromarray(restored_pixels)
    restored_img.save(output_path, "PNG")
    print(f"图像恢复成功，已保存到：{output_path}")
    print(f"最终输出尺寸: {orig_width}x{orig_height}")