"""Sequence-native ComfyUI adapter for the official Apache-2.0 LDF-VFI runtime."""

from __future__ import annotations

import importlib
import logging
from pathlib import Path
import sys
import threading
import types

from einops import rearrange, repeat
import torch
import torch.nn.functional as F


logger = logging.getLogger("Tween")
_LDF_NAMESPACE = "_tween_ldf_upstream"
_LDF_IMPORT_LOCK = threading.RLock()


def _cuda_bf16_supported(device: torch.device) -> bool:
    if device.type != "cuda":
        return False
    with torch.cuda.device(device):
        return torch.cuda.is_bf16_supported()


def _namespace_package(name: str, path: Path):
    package = sys.modules.get(name)
    if package is not None:
        return package
    package = types.ModuleType(name)
    package.__path__ = [str(path)]
    package.__package__ = name
    sys.modules[name] = package
    return package


def load_ldf_runtime(source_root: str):
    """Load LDF under an isolated namespace without polluting ``training``."""
    root = Path(source_root).resolve()
    required = root / "training" / "models" / "precond.py"
    if not required.is_file():
        raise RuntimeError(f"Invalid LDF-VFI source directory: missing {required}")

    with _LDF_IMPORT_LOCK:
        cached = sys.modules.get(f"{_LDF_NAMESPACE}.models.precond")
        if cached is not None:
            transformer = importlib.import_module(f"{_LDF_NAMESPACE}.models.transformer_wan")
            return {
                "Precond": cached.Precond,
                "ConditionalVAE": cached.Wan2_1SpatialTiledConditionEncoder3Dv2,
                "MaskEncoder": cached.MaskSpatialTiledEncoder3D,
                "Transformer": transformer.WanTransformer3DModel,
            }

        training_root = _namespace_package(_LDF_NAMESPACE, root / "training")
        saved_training_modules = {
            name: module for name, module in tuple(sys.modules.items())
            if name == "training" or name.startswith("training.")
        }
        for name in saved_training_modules:
            sys.modules.pop(name, None)
        # One upstream transformer import is absolute (training.distributed.util).
        # Temporarily alias only while importing, then restore the host process.
        sys.modules["training"] = training_root
        try:
            precond = importlib.import_module(f"{_LDF_NAMESPACE}.models.precond")
            transformer = importlib.import_module(f"{_LDF_NAMESPACE}.models.transformer_wan")
        except (ImportError, AttributeError) as exc:
            for name in tuple(sys.modules):
                if name == _LDF_NAMESPACE or name.startswith(f"{_LDF_NAMESPACE}."):
                    sys.modules.pop(name, None)
            raise RuntimeError(
                "LDF-VFI requires PyTorch 2.5+ and diffusers 0.33+. "
                "Install Tween's current requirements and restart ComfyUI."
            ) from exc
        finally:
            for name in tuple(sys.modules):
                if name == "training" or name.startswith("training."):
                    sys.modules.pop(name, None)
            sys.modules.update(saved_training_modules)

    return {
        "Precond": precond.Precond,
        "ConditionalVAE": precond.Wan2_1SpatialTiledConditionEncoder3Dv2,
        "MaskEncoder": precond.MaskSpatialTiledEncoder3D,
        "Transformer": transformer.WanTransformer3DModel,
    }


def _upsample_nearest(frames: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """Expand sparse source frames to every temporal position in a window."""
    kept_indices = torch.where(mask)[0]
    if kept_indices.numel() == 0:
        raise RuntimeError("LDF-VFI received a temporal window with no source frames")
    positions = torch.arange(mask.shape[0], device=mask.device)
    nearest = (positions[:, None] - kept_indices[None, :]).abs().argmin(dim=1)
    return frames[nearest]


class LDFVFIModel:
    """Long-sequence diffusion interpolation using LDF's skip-concat sampler."""

    TRAIN_FRAMES = 60
    TILE_TIME = 20
    CONDITION_TILES = 1

    def __init__(
        self,
        model_root: str,
        vae_path: str,
        source_root: str,
        tile_size: int = 256,
        tile_overlap: int = 64,
        vae_batch_size: int = 8,
        attention_type: str = "slide_chunk_all_block_2x1x1",
    ):
        if tile_size % 8 or tile_overlap % 8:
            raise ValueError("LDF-VFI tile size and overlap must be divisible by 8")
        if tile_overlap >= tile_size:
            raise ValueError("LDF-VFI tile overlap must be smaller than tile size")

        runtime = load_ldf_runtime(source_root)
        self.device = "cpu"
        self.dtype = torch.bfloat16
        self.vae_path = vae_path
        self._ConditionalVAE = runtime["ConditionalVAE"]
        self._MaskEncoder = runtime["MaskEncoder"]
        self._Precond = runtime["Precond"]

        logger.info("Loading LDF-VFI transformer from %s", model_root)
        try:
            transformer = runtime["Transformer"].from_pretrained(
                model_root,
                subfolder="transformer",
                torch_dtype=self.dtype,
                low_cpu_mem_usage=True,
            )
        except TypeError:
            transformer = runtime["Transformer"].from_pretrained(
                model_root, subfolder="transformer", torch_dtype=self.dtype
            )
        transformer.set_attention_type(attention_type)
        transformer.requires_grad_(False).eval()

        stride = tile_size - tile_overlap
        tiled_kwargs = {
            "tile_sample_min_height": tile_size,
            "tile_sample_min_width": tile_size,
            "tile_sample_min_time": self.TILE_TIME,
            "tile_sample_stride_height": stride,
            "tile_sample_stride_width": stride,
            "spatial_compression_ratio": 8,
            "temporal_compression_ratio": 4,
        }
        # The same conditional VAE can encode conditions and decode predictions;
        # sharing it avoids loading a second ~800 MB copy as the reference CLI does.
        self.vae = self._ConditionalVAE(vae_path, vae_batch_size, **tiled_kwargs)
        self.mask_encoder = self._MaskEncoder(**tiled_kwargs)
        self.model = self._Precond(
            transformer=transformer,
            vae=self.vae,
            lq_encoder=self.vae,
            msk_encoder=self.mask_encoder,
        )
        self.model.requires_grad_(False).eval()

    @property
    def transformer(self):
        return self.model.transformer

    def _move_auxiliary_models(self, device: torch.device) -> None:
        self.mask_encoder.mask_encoder.to(device=device, dtype=self.dtype)
        if self.vae.vae is None:
            if device.type != "cpu":
                self.vae.init(device)
            return
        self.vae.vae.to(device=device, dtype=self.dtype)
        if hasattr(self.vae, "mean") and hasattr(self.vae, "std"):
            self.vae.mean = self.vae.mean.to(device=device, dtype=self.dtype)
            self.vae.std = self.vae.std.to(device=device, dtype=self.dtype)
            self.vae.scale = [self.vae.mean, 1.0 / self.vae.std]

    def to(self, device):
        target = torch.device(device)
        if target.type == "cuda" and not _cuda_bf16_supported(target):
            raise RuntimeError("LDF-VFI requires a CUDA GPU with BF16 support (Ampere or newer)")
        self.transformer.to(device=target, dtype=self.dtype)
        self._move_auxiliary_models(target)
        self.device = str(target)
        return self

    def clear_cache(self) -> None:
        for name in (
            "_swin_attention_mask",
            "_sliding_chunk_attention_mask",
            "_sliding_window_attention_mask",
        ):
            method = getattr(self.transformer, name, None)
            cache_clear = getattr(method, "cache_clear", None)
            if cache_clear is not None:
                cache_clear()

    def _prepare_condition(self, frames, mask, device):
        if mask.shape[0] > self.TRAIN_FRAMES:
            raise ValueError("Internal LDF temporal window exceeds the training window")
        if mask.shape[0] < self.TRAIN_FRAMES:
            mask = F.pad(mask, (0, self.TRAIN_FRAMES - mask.shape[0]))
        dense = _upsample_nearest(frames, mask)
        dense = rearrange(dense, "t c h w -> 1 c t h w")
        dense = dense.to(device=device, dtype=self.dtype, non_blocking=True).mul(2).sub(1)
        dense_mask = repeat(
            mask, "t -> 1 1 t h w", h=dense.shape[-2], w=dense.shape[-1]
        ).to(device=device, dtype=self.dtype)
        condition = self.vae.encode(dense, for_train=True)
        encoded_mask = self.mask_encoder.encode(dense_mask, for_train=True)
        return dense, dense_mask, condition, encoded_mask

    def _time_schedule(self, num_steps: int, t_shift: float) -> torch.Tensor:
        schedule = torch.linspace(1.0, 0.0, steps=num_steps + 1)
        return t_shift * schedule / (1 + (t_shift - 1) * schedule)

    def _predict_step(self, latent, timestep, condition, encoded_mask):
        return self.model.predict_v(latent, timestep, condition, encoded_mask)

    def _sample_free(self, condition, encoded_mask, schedule, device, progress):
        latent = torch.randn_like(condition)
        for index in range(schedule.shape[0] - 1):
            progress()
            timestep = torch.full(
                condition.shape[:-4], float(schedule[index]), device=device, dtype=self.dtype
            )
            velocity = self._predict_step(latent, timestep, condition, encoded_mask)
            step_size = float(schedule[index + 1] - schedule[index])
            latent = latent + velocity * step_size
        return latent

    def _sample_between(self, previous, following, condition, encoded_mask,
                        schedule, t_cond, device, progress):
        previous_noisy = previous * (1 - t_cond) + torch.randn_like(previous) * t_cond
        following_noisy = following * (1 - t_cond) + torch.randn_like(following) * t_cond
        middle = torch.randn_like(condition[:, self.CONDITION_TILES:-self.CONDITION_TILES])
        previous_t = torch.full(
            previous.shape[:-4], t_cond, device=device, dtype=self.dtype
        )
        following_t = torch.full(
            following.shape[:-4], t_cond, device=device, dtype=self.dtype
        )
        for index in range(schedule.shape[0] - 1):
            progress()
            latent = torch.cat((previous_noisy, middle, following_noisy), dim=1)
            middle_t = torch.full(
                middle.shape[:-4], float(schedule[index]), device=device, dtype=self.dtype
            )
            timestep = torch.cat((previous_t, middle_t, following_t), dim=1)
            velocity = self._predict_step(
                latent, timestep, condition, encoded_mask
            )[:, self.CONDITION_TILES:-self.CONDITION_TILES]
            step_size = float(schedule[index + 1] - schedule[index])
            middle = middle + velocity * step_size
        return middle

    def _sample_tail(self, previous, condition, encoded_mask, schedule,
                     t_cond, device, progress):
        previous_noisy = previous * (1 - t_cond) + torch.randn_like(previous) * t_cond
        tail = torch.randn_like(condition[:, self.CONDITION_TILES:])
        previous_t = torch.full(
            previous.shape[:-4], t_cond, device=device, dtype=self.dtype
        )
        for index in range(schedule.shape[0] - 1):
            progress()
            latent = torch.cat((previous_noisy, tail), dim=1)
            tail_t = torch.full(
                tail.shape[:-4], float(schedule[index]), device=device, dtype=self.dtype
            )
            timestep = torch.cat((previous_t, tail_t), dim=1)
            velocity = self._predict_step(
                latent, timestep, condition, encoded_mask
            )[:, self.CONDITION_TILES:]
            step_size = float(schedule[index + 1] - schedule[index])
            tail = tail + velocity * step_size
        return tail

    def _decode(self, latent, dense, dense_mask, height, width):
        latent = rearrange(
            latent, "1 nt nh nw c t h w -> 1 nt c t (nh h) (nw w)"
        )
        prediction = self.vae.decode(latent, dense, dense_mask)[..., :height, :width]
        return rearrange(prediction, "1 c t h w -> t c h w").add(1).mul(0.5).clamp_(0, 1).float().cpu()

    @staticmethod
    def sampling_block_count(num_input_frames: int, temporal_factor: int) -> int:
        total_length = num_input_frames * temporal_factor
        t0 = 40
        stride = 20
        blocks = 1
        while t0 + stride <= total_length - 1:
            blocks += 2
            t0 += stride * 2
        if t0 < total_length:
            blocks += 1
        return blocks

    @torch.no_grad()
    def interpolate_sequence(
        self,
        frames: torch.Tensor,
        temporal_factor: int,
        num_steps: int = 16,
        t_shift: float = 8.0,
        t_cond: float = 0.1,
        seed: int = 42,
        progress_callback=None,
    ) -> torch.Tensor:
        if not 2 <= temporal_factor <= 16:
            raise ValueError("LDF-VFI temporal factor must be between 2 and 16")
        if num_steps < 1:
            raise ValueError("LDF-VFI num_steps must be at least 1")
        if t_shift <= 0:
            raise ValueError("LDF-VFI t_shift must be greater than 0")
        if not 0 <= t_cond <= 1:
            raise ValueError("LDF-VFI t_cond must be between 0 and 1")
        if frames.shape[0] < 2:
            return frames

        device = next(self.transformer.parameters()).device
        if device.type != "cuda":
            raise RuntimeError("Move LDF-VFI to a CUDA device before interpolation")
        source = frames.detach().float().cpu()
        height, width = source.shape[-2:]
        schedule = self._time_schedule(num_steps, t_shift)
        progress_callback = progress_callback or (lambda: None)

        cuda_index = device.index
        if cuda_index is None:
            cuda_index = torch.cuda.current_device()
        rng_context = torch.random.fork_rng(devices=[cuda_index])
        with rng_context:
            torch.manual_seed(int(seed))
            with torch.cuda.device(device):
                torch.cuda.manual_seed(int(seed))
            chunks = self._interpolate_sequence_impl(
                source, temporal_factor, schedule, t_cond, device, progress_callback
            )

        result = torch.cat(chunks, dim=0)
        expected = (frames.shape[0] - 1) * temporal_factor + 1
        return result[:expected]

    def _interpolate_sequence_impl(self, source, factor, schedule, t_cond,
                                   device, progress):
        total_mask = torch.zeros(source.shape[0] * factor, dtype=torch.bool)
        total_mask[::factor] = True
        n_tiles = self.TRAIN_FRAMES // self.TILE_TIME
        output_tiles = n_tiles - self.CONDITION_TILES
        outputs = []

        # First chunk.
        mask = total_mask[:self.TRAIN_FRAMES]
        input_end = int(mask.sum())
        dense, dense_mask, condition, encoded_mask = self._prepare_condition(
            source[:input_end], mask, device
        )
        latent = self._sample_free(condition, encoded_mask, schedule, device, progress)
        latent = latent[:, :output_tiles]
        previous = latent[:, -self.CONDITION_TILES:]
        decode_time = output_tiles * self.TILE_TIME
        outputs.append(self._decode(
            latent, dense[:, :, :decode_time], dense_mask[:, :, :decode_time],
            source.shape[-2], source.shape[-1],
        ))

        t0 = output_tiles * self.TILE_TIME
        stride = (n_tiles - self.CONDITION_TILES * 2) * self.TILE_TIME
        while t0 + stride <= total_mask.shape[0] - 1:
            # A future/skip chunk establishes the far-side condition.
            input_start = t0 + stride - self.CONDITION_TILES * self.TILE_TIME
            input_end_t = t0 + stride * 2 + self.CONDITION_TILES * self.TILE_TIME
            mask = total_mask[input_start:input_end_t]
            source_start = int(total_mask[:input_start].sum())
            source_end = int(total_mask[:input_end_t].sum())
            dense, dense_mask, condition, encoded_mask = self._prepare_condition(
                source[source_start:source_end], mask, device
            )
            skip = self._sample_free(condition, encoded_mask, schedule, device, progress)
            skip = skip[:, self.CONDITION_TILES:-self.CONDITION_TILES]
            following = skip[:, :self.CONDITION_TILES]
            previous_next = skip[:, -self.CONDITION_TILES:]
            start_time = self.CONDITION_TILES * self.TILE_TIME
            end_time = output_tiles * self.TILE_TIME
            decoded_skip = self._decode(
                skip, dense[:, :, start_time:end_time], dense_mask[:, :, start_time:end_time],
                source.shape[-2], source.shape[-1],
            )

            # Fill the gap between the preceding and skip chunks.
            input_start = t0 - self.CONDITION_TILES * self.TILE_TIME
            input_end_t = t0 + stride + self.CONDITION_TILES * self.TILE_TIME
            mask = total_mask[input_start:input_end_t]
            source_start = int(total_mask[:input_start].sum())
            source_end = int(total_mask[:input_end_t].sum())
            dense, dense_mask, condition, encoded_mask = self._prepare_condition(
                source[source_start:source_end], mask, device
            )
            middle = self._sample_between(
                previous, following, condition, encoded_mask,
                schedule, t_cond, device, progress,
            )
            decoded_middle = self._decode(
                middle, dense[:, :, start_time:end_time], dense_mask[:, :, start_time:end_time],
                source.shape[-2], source.shape[-1],
            )
            outputs.extend((decoded_middle, decoded_skip))
            previous = previous_next
            t0 += stride * 2

        # Remaining tail.
        if t0 < total_mask.shape[0]:
            input_start = t0 - self.CONDITION_TILES * self.TILE_TIME
            mask = total_mask[input_start:]
            source_start = int(total_mask[:input_start].sum())
            dense, dense_mask, condition, encoded_mask = self._prepare_condition(
                source[source_start:], mask, device
            )
            tail = self._sample_tail(
                previous, condition, encoded_mask, schedule,
                t_cond, device, progress,
            )
            start_time = self.CONDITION_TILES * self.TILE_TIME
            outputs.append(self._decode(
                tail, dense[:, :, start_time:], dense_mask[:, :, start_time:],
                source.shape[-2], source.shape[-1],
            ))

        return outputs