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

from moviepy import VideoFileClip
from .pixel_scramble import obfuscate, deobfuscate  # 根据实际路径调整

def _process_frame(frame: np.ndarray, func, seed: Optional[int], use_timestamp: bool) -> np.ndarray:
    """辅助函数：对单个帧应用混淆/恢复函数"""
    pil_img = Image.fromarray(frame)

    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:
        pil_img.save(tmp_in_path, 'PNG')
        if func == obfuscate:
            func(tmp_in_path, tmp_out_path, seed=seed, use_timestamp=use_timestamp)
        else:
            func(tmp_in_path, tmp_out_path, seed=seed)
        result_pil = Image.open(tmp_out_path)
        result_array = np.array(result_pil)
        if result_array.shape[-1] == 4:
            result_array = result_array[..., :3]  # 丢弃 alpha 通道
        return result_array
    finally:
        os.remove(tmp_in_path)
        os.remove(tmp_out_path)

def obfuscate_mp4(
    input_path: str,
    output_path: str,
    seed: Optional[int] = None,
    use_timestamp: bool = False
) -> None:
    """对 MP4 视频逐帧混淆（加密）"""
    if use_timestamp:
        warnings.warn("use_timestamp will be ignored for mp4 files")
        use_timestamp = False

    clip = VideoFileClip(input_path)
    print(f"视频信息：时长={clip.duration}s，帧率={clip.fps}")

    # transform 要求函数接收 get_frame 和 t，我们需要手动调用 get_frame(t) 获取帧
    def process_frame(get_frame, t):
        frame = get_frame(t)  # 获取当前帧的 numpy 数组
        return _process_frame(frame, obfuscate, seed, use_timestamp)

    processed_clip = clip.transform(process_frame)
    processed_clip.write_videofile(
        output_path,
        codec='libx264rgb',  # RGB 无损编码
        ffmpeg_params=["-crf", "0", "-preset", "ultrafast"],  # crf=0 为无损
        audio_codec='aac'  # 音频仍可用有损，不影响视频
    )
    print(f"MP4 混淆完成，已保存到：{output_path}")

def deobfuscate_mp4(
    input_path: str,
    output_path: str,
    seed: Optional[int] = None
) -> None:
    """恢复被混淆的 MP4 视频"""
    clip = VideoFileClip(input_path)

    def process_frame(get_frame, t):
        frame = get_frame(t)
        return _process_frame(frame, deobfuscate, seed, False)

    processed_clip = clip.transform(process_frame)
    processed_clip.write_videofile(
        output_path,
        codec='libx264rgb',
        ffmpeg_params=["-crf", "0", "-preset", "ultrafast"],  # crf=0 为无损
        audio_codec='aac'  # 音频仍可用有损，不影响视频
    )
    print(f"MP4 恢复完成，已保存到：{output_path}")