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

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)  # 强制在 0 ~ 4294967295 内


def obfuscate(
    image_path: str,
    output_path: str,
    seed: Optional[int] = None,
    use_timestamp: bool = True
) -> Optional[int]:
    img = Image.open(image_path)
    pixels = np.array(img)
    original_shape = pixels.shape
    flat_pixels = pixels.flatten()

    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_pixels))
    rng.shuffle(indices)

    shuffled_flat = flat_pixels[indices]
    shuffled_pixels = shuffled_flat.reshape(original_shape)

    shuffled_img = Image.fromarray(shuffled_pixels)

    # 保存 PNG 元数据
    info = PngImagePlugin.PngInfo()
    if timestamp is not None:
        info.add_text("obf_timestamp", str(timestamp))
        print(f"已将 timestamp {timestamp} 嵌入图像元数据")

    shuffled_img.save(output_path, "PNG", pnginfo=info)
    print(f"图像混淆成功，已保存到：{output_path}")

    return timestamp


def deobfuscate(
    obf_image_path: str,
    output_path: str,
    seed: Optional[int] = None,
    timestamp: Optional[int] = None
) -> None:

    img = Image.open(obf_image_path)
    pixels = np.array(img)
    original_shape = pixels.shape
    flat_pixels = pixels.flatten()

    # 尝试从元数据读取 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 格式无效，将忽略")

    # 决定最终使用的 timestamp
    final_timestamp = ts_from_meta if ts_from_meta is not None else timestamp

    # 如果最终 timestamp 为 None，则使用纯 seed 模式
    if final_timestamp is None:
        print("未找到 timestamp，使用纯 seed 模式进行恢复（假设混淆时未启用 timestamp）")

    final_seed = _generate_final_seed(seed, final_timestamp)

    rng = np.random.RandomState(final_seed)
    indices = np.arange(len(flat_pixels))
    rng.shuffle(indices)

    inverse_indices = np.argsort(indices)
    restored_flat = flat_pixels[inverse_indices]

    restored_pixels = restored_flat.reshape(original_shape)
    restored_img = Image.fromarray(restored_pixels)
    restored_img.save(output_path, "PNG")
    print(f"图像恢复成功，已保存到：{output_path}")