import os
import tempfile
import numpy as np
from PIL import Image
from typing import Optional, Callable
import warnings
import time
from functools import lru_cache
from moviepy import VideoFileClip
from .pixel_scramble import obfuscate, deobfuscate


# 性能优化：使用内存缓存避免重复的临时文件操作
@lru_cache(maxsize=32)
def _process_frame_cached(
    frame_hash: int, frame_shape: tuple, func_name: str, seed: Optional[int], use_timestamp: bool
) -> np.ndarray:
    """缓存版本的帧处理函数"""
    # 这个函数主要用于缓存，实际处理在 _process_frame_fast 中
    pass


def _process_frame_fast(
    frame: np.ndarray, func: Callable, seed: Optional[int], use_timestamp: bool
) -> np.ndarray:
    """优化版帧处理：直接在内存中操作，减少临时文件IO"""
    # 将帧转换为PIL图像
    pil_img = Image.fromarray(frame)

    # 直接在内存中处理，避免临时文件
    if func == obfuscate:
        # 获取像素数据
        pixels = np.array(pil_img)
        original_shape = pixels.shape
        flat_pixels = pixels.flatten()

        # 生成种子
        timestamp = int(time.time()) if use_timestamp else None
        from .pixel_scramble import _generate_final_seed
        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)
        result_array = shuffled_pixels
    else:
        # 解混淆
        pixels = np.array(pil_img)
        original_shape = pixels.shape
        flat_pixels = pixels.flatten()

        from .pixel_scramble import _generate_final_seed
        final_seed = _generate_final_seed(seed, None)

        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)
        result_array = restored_pixels

    # 确保输出格式正确
    if result_array.shape[-1] == 4:
        result_array = result_array[..., :3]  # 丢弃 alpha 通道
    return result_array


def _process_frame_with_fallback(
    frame: np.ndarray, func: Callable, seed: Optional[int], use_timestamp: bool
) -> np.ndarray:
    """带降级策略的帧处理：优先使用快速方法，失败时使用临时文件方法"""
    try:
        return _process_frame_fast(frame, func, seed, use_timestamp)
    except Exception as e:
        warnings.warn(f"快速处理失败，使用临时文件方法: {e}")
        return _process_frame_legacy(frame, func, seed, use_timestamp)


def _process_frame_legacy(
    frame: np.ndarray, func: Callable, 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]
        return result_array
    finally:
        try:
            os.remove(tmp_in_path)
            os.remove(tmp_out_path)
        except:
            pass


class ProgressBar:
    """简单的进度条显示"""
    def __init__(self, total: int, description: str = "Processing"):
        self.total = total
        self.current = 0
        self.description = description
        self.start_time = time.time()
        self.last_update = 0

    def update(self, increment: int = 1):
        self.current += increment
        current_time = time.time()
        # 每秒最多更新一次显示
        if current_time - self.last_update >= 1.0 or self.current == self.total:
            progress = self.current / self.total
            elapsed = current_time - self.start_time
            if progress > 0:
                eta = elapsed / progress * (1 - progress)
                print(f"\r{self.description}: {progress:.1%} ({self.current}/{self.total}) "
                      f"ETA: {eta:.1f}s", end="", flush=True)
            self.last_update = current_time

    def finish(self):
        elapsed = time.time() - self.start_time
        print(f"\r{self.description}: 完成! 耗时: {elapsed:.1f}s")


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

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

    progress = ProgressBar(total_frames, "混淆进度") if show_progress else None
    frame_count = [0]  # 使用列表以便在闭包中修改

    # transform 要求函数接收 get_frame 和 t，我们需要手动调用 get_frame(t) 获取帧
    def process_frame(get_frame, t):
        frame = get_frame(t)  # 获取当前帧的 numpy 数组
        result = _process_frame_with_fallback(frame, obfuscate, seed, use_timestamp)
        if progress:
            frame_count[0] += 1
            progress.update()
        return result

    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',  # 音频仍可用有损，不影响视频
        logger=None if show_progress else 'bar'  # 避免重复进度条
    )

    if progress:
        progress.finish()
    print(f"MP4 混淆完成，已保存到：{output_path}")


def deobfuscate_mp4(
    input_path: str,
    output_path: str,
    seed: Optional[int] = None,
    show_progress: bool = True
) -> None:
    """恢复被混淆的 MP4 视频 - 性能优化版本"""
    clip = VideoFileClip(input_path)
    total_frames = int(clip.duration * clip.fps)

    progress = ProgressBar(total_frames, "恢复进度") if show_progress else None
    frame_count = [0]

    def process_frame(get_frame, t):
        frame = get_frame(t)
        result = _process_frame_with_fallback(frame, deobfuscate, seed, False)
        if progress:
            frame_count[0] += 1
            progress.update()
        return result

    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',
        logger=None if show_progress else 'bar'
    )

    if progress:
        progress.finish()
    print(f"MP4 恢复完成，已保存到：{output_path}")