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 = [str(seed) if seed is not None else "default"]
    if timestamp is not None:
        parts.append(str(timestamp))
    combined = "_".join(parts)
    return int(hashlib.sha256(combined.encode()).hexdigest(), 16) % (2**32)


def _pad_to_square(img: Image.Image) -> Tuple[Image.Image, int, int, int, int]:
    width, height = img.size
    side = max(width, height)
    delta_w = side - width
    delta_h = side - height
    pad_left = delta_w // 2
    pad_top = delta_h // 2
    pad_right = delta_w - pad_left
    pad_bottom = delta_h - pad_top
    padded_img = Image.new("RGB", (side, side), (128, 128, 128))
    padded_img.paste(img, (pad_left, pad_top))
    return padded_img, width, height, pad_left, pad_top


def obfuscate(
    image_path: str,
    output_path: str,
    iterations: int = 50,
    seed: Optional[int] = None,
    use_timestamp: bool = True
) -> Optional[int]:
    img = Image.open(image_path).convert("RGB")
    padded_img, orig_width, orig_height, pad_left, pad_top = _pad_to_square(img)
    n = padded_img.size[0]

    pixels = np.array(padded_img)

    timestamp = int(time.time()) if use_timestamp else None
    final_seed = _generate_final_seed(seed, timestamp)

    # 正向迭代
    for _ in range(iterations):
        new_pixels = np.zeros_like(pixels)
        for y in range(n):
            for x in range(n):
                new_x = (2 * x + y) % n
                new_y = (x + y) % n
                new_pixels[new_y, new_x] = pixels[y, x]
        pixels = new_pixels

    scrambled_img = Image.fromarray(pixels)

    info = PngImagePlugin.PngInfo()
    if timestamp is not None:
        info.add_text("arnold_timestamp", str(timestamp))
    info.add_text("arnold_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("pad_left", str(pad_left))
    info.add_text("pad_top", str(pad_top))

    scrambled_img.save(output_path, "PNG", pnginfo=info)
    print(f"Arnold 猫映射混淆完成，已保存到：{output_path}")
    print(f"迭代次数: {iterations}，方形边长: {n}")
    if timestamp is not None:
        print(f"timestamp 已嵌入: {timestamp}")

    return timestamp


def deobfuscate(
    obf_image_path: str,
    output_path: str,
    iterations: Optional[int] = None,
    seed: Optional[int] = None,
    timestamp: Optional[int] = None
) -> None:
    img = Image.open(obf_image_path).convert("RGB")
    pixels = np.array(img)
    n = pixels.shape[0]

    # 读取元数据
    ts_from_meta = None
    ts_str = img.info.get("arnold_timestamp")
    if ts_str:
        ts_from_meta = int(ts_str)

    iters_from_meta = None
    iters_str = img.info.get("arnold_iterations")
    if iters_str:
        iters_from_meta = int(iters_str)

    final_timestamp = ts_from_meta if ts_from_meta is not None else timestamp
    final_iterations = iters_from_meta if iters_from_meta is not None else iterations

    if final_iterations is None:
        raise ValueError("必须提供 iterations 或从元数据读取")

    orig_width = int(img.info.get("orig_width", n))
    orig_height = int(img.info.get("orig_height", n))
    padded_size = int(img.info.get("padded_size", n))
    pad_left = int(img.info.get("pad_left", 0))
    pad_top = int(img.info.get("pad_top", 0))

    print(f"恢复参数：iterations={final_iterations}, timestamp={final_timestamp}")
    print(f"原始尺寸: {orig_width}x{orig_height}, padded: {padded_size}")
    print(f"填充位置: pad_left={pad_left}, pad_top={pad_top}")

    # 逆向迭代
    for _ in range(final_iterations):
        new_pixels = np.zeros_like(pixels)
        for y in range(n):
            for x in range(n):
                prev_x = (x - y) % n
                prev_y = (-x + 2 * y) % n
                new_pixels[prev_y, prev_x] = pixels[y, x]
        pixels = new_pixels

    # 裁剪回原始尺寸（从填充位置开始裁剪）
    restored_pixels = pixels[pad_top:pad_top + orig_height, pad_left:pad_left + orig_width, :]
    restored_img = Image.fromarray(restored_pixels)
    restored_img.save(output_path, "PNG")
    print(f"Arnold 猫映射恢复完成，已保存到：{output_path}")
    print(f"最终尺寸: {orig_width}x{orig_height}")