import gc
import os
import torch
import pathlib
import einops
import folder_paths
import typing
import time
import urllib.request
import urllib.error
#from vfi_utils import load_file_from_github_release
from comfy.model_management import get_torch_device, soft_empty_cache
from packaging import version
from .rife_model import IFNet

MODEL_TYPE = pathlib.Path(__file__).parent.name

CKPT_NAME_VER_DICT = {
    "rife47.pth": "4.7",
    "rife49.pth": "4.7",
    "rife417.pth": "4.17",
}

CKPT_NAMES = list(CKPT_NAME_VER_DICT.keys())

DTYPE_OPTIONS = ["float32", "float16", "bfloat16"]
DTYPE_MAP = {
    "float32": torch.float32,
    "float16": torch.float16,
    "bfloat16": torch.bfloat16,
}

# Module-level model cache: avoids reloading weights on every node execution.
# Key: (ckpt_name, dtype, torch_compile) — invalidated if any of these change.
_model_cache: typing.Dict[typing.Tuple, torch.nn.Module] = {}


def download_model(ckpt_name: str, model_path: str) -> bool:
    """Download model from GitHub releases if it doesn't exist locally."""
    if os.path.exists(model_path):
        return True

    if ckpt_name not in CKPT_NAME_VER_DICT:
        print(f"Comfy-VFI: Unknown model '{ckpt_name}', cannot download.")
        return False

    url = f"https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/releases/download/models/{ckpt_name}"
    os.makedirs(os.path.dirname(model_path), exist_ok=True)

    try:
        print(f"Comfy-VFI: Downloading {ckpt_name} from {url}...")
        urllib.request.urlretrieve(url, model_path)
        print(f"Comfy-VFI: Successfully downloaded {ckpt_name}")
        return True
    except urllib.error.URLError as e:
        print(f"Comfy-VFI: Failed to download {ckpt_name}: {e}")
        return False
    except Exception as e:
        print(f"Comfy-VFI: Unexpected error downloading {ckpt_name}: {e}")
        return False


class InterpolationStateList():

    def __init__(self, frame_indices: typing.List[int], is_skip_list: bool):
        self.frame_indices = frame_indices
        self.is_skip_list = is_skip_list
        
    def is_frame_skipped(self, frame_index):
        is_frame_in_list = frame_index in self.frame_indices
        return self.is_skip_list and is_frame_in_list or not self.is_skip_list and not is_frame_in_list
    

def preprocess_frames(frames):
    return einops.rearrange(frames[..., :3], "n h w c -> n c h w")

def postprocess_frames(frames):
    return einops.rearrange(frames, "n c h w -> n h w c")[..., :3].cpu()


class RIFE_VFI:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "ckpt_name": (CKPT_NAMES,),
                "frames": ("IMAGE", ),
                "clear_cache_after_n_frames": ("INT", {"default": 10, "min": 1, "max": 1000}),
                "multiplier": ("INT", {"default": 2, "min": 1}),
                "fast_mode": ("BOOLEAN", {"default": True}),
                "ensemble": ("BOOLEAN", {"default": True}),
                "scale_factor": ([0.25, 0.5, 1.0, 2.0, 4.0], {"default": 1.0}),
                # Opt 1: Precision — float32 is safe, float16 is ~2x faster/half VRAM,
                #         bfloat16 is safer than float16 on Ampere+ GPUs (RTX 30xx+)
                "dtype": (DTYPE_OPTIONS, {"default": "bfloat16"}),
                # Opt 2: torch.compile() — JIT-compiles the model graph for 10-30% speedup.
                #         First inference will be slow (compilation), subsequent ones faster.
                #         Requires PyTorch 2.0+. Disable if you hit errors on older builds.
                "torch_compile": ("BOOLEAN", {"default": False,
                                              "tooltip": "Compile the model with torch.compile() for 10-30% faster inference "
                                                         "after the first (warm-up) run. Requires PyTorch 2.0+."}),
                # Opt 7: Batch tasks — process multiple (pair, timestep) tasks per GPU call.
                #         Each task is one interpolated frame. Batching amortises kernel launch
                #         overhead; higher values improve throughput but use more VRAM.
                #         1 = one interpolated frame per call (safest).
                "batch_size": ("INT", {"default": 1, "min": 1, "max": 64,
                                       "tooltip": "Number of interpolation tasks per GPU call. "
                                                  "Higher values improve throughput but use more VRAM. "
                                                  "Set to 1 for the most conservative behaviour."}),
            },
            "optional": {
                "optional_interpolation_states": ("INTERPOLATION_STATES", )
            }
        }

    RETURN_TYPES = ("IMAGE", )
    FUNCTION = "vfi"
    CATEGORY = "ComfyUI-Frame-Interpolation/VFI"

    def _get_model_path(self, model_name: str) -> str:
        """获取模型文件的完整路径"""
        rife_models_folder = os.path.join(folder_paths.models_dir, "rife")
        return os.path.join(rife_models_folder, model_name)

    def vfi(
        self,
        ckpt_name: typing.AnyStr,
        frames: torch.Tensor,
        clear_cache_after_n_frames: int = 10,
        multiplier: typing.SupportsInt = 2,
        fast_mode: bool = True,
        ensemble: bool = True,
        scale_factor: float = 1.0,
        dtype: str = "bfloat16",
        torch_compile: bool = True,
        batch_size: int = 1,
        optional_interpolation_states: InterpolationStateList = None,
        **kwargs
    ):
        """
        Perform video frame interpolation using a given checkpoint model.

        Args:
            ckpt_name (str): The name of the checkpoint model to use.
            frames (torch.Tensor): A tensor containing input video frames.
            clear_cache_after_n_frames (int, optional): The number of frames to process before clearing CUDA cache
                to prevent memory overflow. Defaults to 10. Lower numbers are safer but mean more processing time.
                How high you should set it depends on how many input frames there are, input resolution (after upscaling),
                how many times you want to multiply them, and how long you're willing to wait for the process to complete.
            multiplier (int, optional): The multiplier for each input frame. 60 input frames * 2 = 120 output frames. Defaults to 2.
            dtype (str, optional): Floating-point precision for inference.
                "float32" — full precision, safest (default).
                "float16" — half precision, ~2x faster and ~50% less VRAM, may have minor artifacts.
                "bfloat16" — brain float16, better numeric range than float16, requires Ampere+ GPU (RTX 30xx+).
            torch_compile (bool, optional): Compile the model with torch.compile() for 10-30% speedup.
                The first inference call will be slow (compilation warm-up); subsequent calls are faster.
                Requires PyTorch 2.0+. Defaults to False.
            batch_size (int, optional): Number of interpolation tasks per GPU call. Each task produces one
                intermediate frame. Higher values improve throughput but use more VRAM. Defaults to 1.

        Returns:
            tuple: A tuple containing the output interpolated frames.
        """
        start_time = time.perf_counter()
        model_path = self._get_model_path(ckpt_name)
        arch_ver = CKPT_NAME_VER_DICT[ckpt_name]

        # Download model if it doesn't exist
        if not download_model(ckpt_name, model_path):
            raise RuntimeError(f"Model {ckpt_name} not found and could not be downloaded.")

        torch_dtype = DTYPE_MAP[dtype]
        device = get_torch_device()
        print(f'{time.time():.3f}: device: {device}, dtype: {frames.dtype}')
        #  4.26 disable ensemble
        if arch_ver == "4.26":
            ensemble = False  # 4.26 版本不支持 ensemble

        # Cache the model by (ckpt_name, dtype, torch_compile) so repeated node
        # executions skip the load_state_dict + device transfer entirely.
        cache_key = (ckpt_name, dtype, torch_compile)
        if cache_key not in _model_cache:
            interpolation_model = IFNet(arch_ver=arch_ver)
            interpolation_model.load_state_dict(torch.load(model_path, map_location=device, weights_only=False))
            if torch_dtype != torch.float32:
                interpolation_model = interpolation_model.to(torch_dtype)
            interpolation_model.eval().to(device)
            # Opt 2: torch.compile() — JIT-compiles the model for 10-30% faster inference.
            # The first call triggers compilation (slow); all subsequent calls use the cached graph.
            if torch_compile:
                interpolation_model = torch.compile(interpolation_model)
            _model_cache[cache_key] = interpolation_model
            print(f"{time.time():.3f}: Comfy-VFI: Loaded and cached model {ckpt_name} ({dtype}{'+ torch.compile' if torch_compile else ''})")
        else:
            interpolation_model = _model_cache[cache_key]
            print(f"{time.time():.3f}: Comfy-VFI: Using cached model {ckpt_name} ({dtype}{'+ torch.compile' if torch_compile else ''})")

        print(
            f"Comfy-VFI: model device={next(interpolation_model.parameters()).device}, "
            f"dtype={next(interpolation_model.parameters()).dtype}"
        )

        frames = preprocess_frames(frames)

        # Normalise multiplier to a per-pair list
        n_pairs = len(frames) - 1
        if isinstance(multiplier, int):
            multipliers = [int(multiplier)] * n_pairs
        else:
            multipliers = list(map(int, multiplier))
            multipliers += [2] * (n_pairs - len(multipliers))

         #  4.26 scale_list
        if arch_ver == "4.26":
            scale_list = [16 / scale_factor, 8 / scale_factor, 4 / scale_factor, 2 / scale_factor, 1 / scale_factor]
        else:
            scale_list = [8 / scale_factor, 4 / scale_factor, 2 / scale_factor, 1 / scale_factor]

        # Build a flat list of all (pair_idx, timestep) tasks, skipping excluded pairs.
        # Each task produces exactly one intermediate frame.
        tasks: typing.List[typing.Tuple[int, float]] = []
        tasks_remaining_per_pair: typing.Dict[int, int] = {}
        for pair_idx in range(n_pairs):
            if optional_interpolation_states is not None and optional_interpolation_states.is_frame_skipped(pair_idx):
                tasks_remaining_per_pair[pair_idx] = 0
                continue
            m = multipliers[pair_idx]
            n_steps = max(m - 1, 0)
            tasks_remaining_per_pair[pair_idx] = n_steps
            for step in range(1, m):
                tasks.append((pair_idx, step / m))

        # Storage for intermediate frames, keyed by pair index
        results: typing.Dict[int, typing.List[torch.Tensor]] = {i: [] for i in range(n_pairs)}

        frames_processed_since_cache_clear = 0
        pos = 0

        # Opt 5: inference_mode disables autograd tracking entirely — lower overhead than no_grad().
        # Opt 7: process batch_size tasks per GPU call. IFNet supports batched tensor timesteps,
        #         so multiple (pair, timestep) combinations can be inferred in a single forward pass.
        print(f"{time.time():.3f}: Comfy-VFI: start inference")
        total_inference_time = 0.0
        total_gpu_to_cpu_time = 0.0
        
        with torch.inference_mode():
            while pos < len(tasks):
                batch_tasks = tasks[pos : pos + batch_size]

                frame0_list, frame1_list, timestep_list = [], [], []
                for pair_idx, dt in batch_tasks:
                    frame0_list.append(frames[pair_idx : pair_idx + 1])
                    frame1_list.append(frames[pair_idx + 1 : pair_idx + 2])
                    timestep_list.append(dt)

                frame0_batch = torch.cat(frame0_list, dim=0).to(device, dtype=torch_dtype)
                frame1_batch = torch.cat(frame1_list, dim=0).to(device, dtype=torch_dtype)
                # Batched timestep tensor: shape [B, 1, 1, 1] — IFNet expands it internally
                timestep_tensor = torch.tensor(timestep_list, dtype=torch_dtype, device=device).view(-1, 1, 1, 1)

                t_infer = time.perf_counter()
                
                middle_frames = interpolation_model(
                    frame0_batch,
                    frame1_batch,
                    timestep_tensor,
                    scale_list,
                    fast_mode,
                    ensemble,
                ).clamp(0, 1)

                infer_elapsed = time.perf_counter() - t_infer
                total_inference_time += infer_elapsed

                t_copy = time.perf_counter()
                middle_frames = middle_frames.detach().cpu()
                copy_elapsed = time.perf_counter() - t_copy
                total_gpu_to_cpu_time += copy_elapsed
                
                for idx, (pair_idx, _) in enumerate(batch_tasks):
                    results[pair_idx].append(middle_frames[idx : idx + 1].to(dtype=torch_dtype))
                    tasks_remaining_per_pair[pair_idx] -= 1
                    # Clear cache after finishing each pair's worth of tasks
                    if tasks_remaining_per_pair[pair_idx] == 0:
                        frames_processed_since_cache_clear += 1
                        if frames_processed_since_cache_clear >= clear_cache_after_n_frames:
                            soft_empty_cache()
                            gc.collect()
                            frames_processed_since_cache_clear = 0

                pos += len(batch_tasks)

        # Assemble output: each original frame followed by its interpolated frames
        output_frames: typing.List[torch.Tensor] = []
        for pair_idx in range(n_pairs):
            output_frames.append(frames[pair_idx : pair_idx + 1].to(dtype=torch_dtype))
            for mid in results[pair_idx]:
                output_frames.append(mid)
        output_frames.append(frames[-1:].to(dtype=torch_dtype))

        print(f"{time.time():.3f}: Comfy-VFI: Final clearing cache...")
        soft_empty_cache()
        print(f"{time.time():.3f}: Done cache clearing")
        print(f"{time.time():.3f}: Comfy-VFI inference_time: {total_inference_time:.4f}; copy_time: {total_gpu_to_cpu_time:.4f}")
        print(f"{time.time():.3f}: Comfy-VFI done! {sum(len(v) for v in results.values()) + len(frames)} frames generated")
        
        # Always return float32 — numpy and all downstream ComfyUI nodes require it
        out_tensor = torch.cat(output_frames, dim=0).to(torch.float32)

        stop_time = time.perf_counter()
        elapsed_time = stop_time - start_time
        print(f"{time.time():.3f}: VFI elapsed: {elapsed_time:.3f} s ({elapsed_time * 1000:.2f} ms)\n")
        return (postprocess_frames(out_tensor),)