import os
import tempfile
import numpy as np
from PIL import Image
from typing import Optional

# 根据实际路径导入原有的混淆/恢复函数
from .pixel_scramble import obfuscate, deobfuscate

def obfuscate_gif(
    gif_path: str,
    output_gif_path: str,
    seed: Optional[int] = None,
    use_timestamp: bool = False
) -> None:
    """
    对 GIF 逐帧混淆（加密），保存为新的 GIF。
    建议 use_timestamp=False，以保证可逆性。
    """
    gif = Image.open(gif_path)
    frames_data = []          # 存储加密后的灰度图像 (L模式)
    durations = []
    palettes = []             # 存储每帧的原始调色板
    transparencies = []       # 存储每帧的透明索引（若有）

    try:
        while True:
            frame = gif.copy()
            durations.append(gif.info.get('duration', 100))

            # 确保帧为 P 模式（索引颜色）
            if frame.mode != 'P':
                frame = frame.convert('P')
            palette = frame.getpalette()
            if palette is None:
                raise ValueError("原始帧缺少调色板信息，无法加密")
            transparency = frame.info.get('transparency')

            # 保存原始调色板和透明信息
            palettes.append(palette)
            transparencies.append(transparency)

            # 获取索引数组 (H, W) dtype=uint8
            indices = np.array(frame)

            # 将索引数组保存为灰度 PNG（临时文件）
            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_in:
                tmp_in_path = tmp_in.name
            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_out:
                tmp_out_path = tmp_out.name

            try:
                idx_img = Image.fromarray(indices.astype(np.uint8), mode='L')
                idx_img.save(tmp_in_path, 'PNG')

                # 调用原有的混淆函数（仅位置重排）
                obfuscate(tmp_in_path, tmp_out_path, seed=seed, use_timestamp=use_timestamp)

                # 读取加密后的灰度图像
                enc_gray = Image.open(tmp_out_path).convert('L')
                frames_data.append(enc_gray)
            finally:
                os.remove(tmp_in_path)
                os.remove(tmp_out_path)

            gif.seek(gif.tell() + 1)
    except EOFError:
        pass  # 所有帧处理完毕

    if not frames_data:
        raise ValueError("GIF 没有有效帧")

    # 将加密后的灰度图像转换为 P 模式，应用原始调色板和透明信息，保存为 GIF
    def to_p_mode(gray_img: Image.Image, pal: list, trans: Optional[int]) -> Image.Image:
        arr = np.array(gray_img).astype(np.uint8)
        p_img = Image.fromarray(arr, mode='P')
        p_img.putpalette(pal)
        if trans is not None:
            p_img.info['transparency'] = trans
        return p_img

    first = to_p_mode(frames_data[0], palettes[0], transparencies[0])
    rest = [to_p_mode(g, p, t) for g, p, t in zip(frames_data[1:], palettes[1:], transparencies[1:])]

    first.save(
        output_gif_path,
        save_all=True,
        append_images=rest,
        duration=durations,
        loop=0,
        disposal=2          # 清理背景，避免透明残留
    )
    print(f"GIF 混淆完成，已保存到：{output_gif_path}")

def deobfuscate_gif(
    obf_gif_path: str,
    output_gif_path: str,
    seed: Optional[int] = None
) -> None:
    """
    恢复被混淆的 GIF。
    假设加密时使用了相同的 seed，且 use_timestamp=False。
    """
    gif = Image.open(obf_gif_path)
    frames = []
    durations = []

    try:
        while True:
            frame = gif.copy()
            durations.append(gif.info.get('duration', 100))

            # 直接从当前帧获取调色板和透明信息（加密时已保留原始信息）
            if frame.mode != 'P':
                frame = frame.convert('P')
            palette = frame.getpalette()
            if palette is None:
                raise ValueError("加密帧缺少调色板信息")
            transparency = frame.info.get('transparency')

            # 获取加密后的索引数组
            indices = np.array(frame)

            # 保存索引数组为灰度 PNG（临时文件）
            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_in:
                tmp_in_path = tmp_in.name
            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_out:
                tmp_out_path = tmp_out.name

            try:
                idx_img = Image.fromarray(indices.astype(np.uint8), mode='L')
                idx_img.save(tmp_in_path, 'PNG')

                # 调用原有的解密函数
                deobfuscate(tmp_in_path, tmp_out_path, seed=seed)

                # 读取恢复的索引数组
                restored_img = Image.open(tmp_out_path)
                restored_indices = np.array(restored_img)

                # 重建 P 模式帧，应用原始调色板和透明信息
                restored_p = Image.fromarray(restored_indices.astype(np.uint8), mode='P')
                restored_p.putpalette(palette)
                if transparency is not None:
                    restored_p.info['transparency'] = transparency

                frames.append(restored_p)
            finally:
                os.remove(tmp_in_path)
                os.remove(tmp_out_path)

            gif.seek(gif.tell() + 1)
    except EOFError:
        pass

    if frames:
        first = frames[0]
        first.save(
            output_gif_path,
            save_all=True,
            append_images=frames[1:],
            duration=durations,
            loop=0,
            disposal=2
        )
    print(f"GIF 恢复完成，已保存到：{output_gif_path}")