from typing import Any, Optional, Tuple

import torch
import torch.nn.functional as F

from .core_sampling import (
    SampleChunkConfig,
    decode_latent_batch_chunked,
    encode_image_batch_chunked,
    sample_latent_batch_chunked,
    sample_single_latent_temporal,
    soft_clear_cache,
)

try:
    import cv2  # type: ignore
except ImportError:  # pragma: no cover
    cv2 = None


def _temporal_rgb_view(image: torch.Tensor) -> torch.Tensor:
    if image.ndim != 4:
        raise ValueError(f"expected IMAGE tensor with shape [B,H,W,C], got {tuple(image.shape)}")
    if image.shape[-1] > 3:
        return image[..., :3]
    return image


def _frame_batch_size(images: torch.Tensor) -> int:
    if images.ndim != 4:
        raise ValueError(f"expected IMAGE tensor with shape [B,H,W,C], got {tuple(images.shape)}")
    return int(images.shape[0])


def _single_frame(images: torch.Tensor, index: int) -> torch.Tensor:
    return images[index : index + 1]


def _validate_optional_frame_batch(name: str, value: Optional[torch.Tensor], batch_size: int) -> None:
    if value is None:
        return
    actual = _frame_batch_size(value)
    if actual not in {1, batch_size, max(1, batch_size - 1)}:
        raise ValueError(
            f"{name} batch size must be 1, the full frame count, or frame_count - 1; "
            f"got {actual} for {batch_size} frames"
        )


def _extract_optional_frame(
    name: str,
    value: Optional[torch.Tensor],
    frame_index: int,
    batch_size: int,
) -> Optional[torch.Tensor]:
    if value is None:
        return None

    actual = _frame_batch_size(value)
    if actual == 1:
        return value
    if actual == batch_size:
        return _single_frame(value, frame_index)
    if actual == batch_size - 1:
        if frame_index == 0:
            return None
        return _single_frame(value, frame_index - 1)
    raise ValueError(f"{name} batch size mismatch")


def _mask_batch_size(mask: torch.Tensor) -> int:
    if mask.ndim == 2:
        return 1
    if mask.ndim == 3:
        return int(mask.shape[0])
    if mask.ndim == 4 and mask.shape[-1] == 1:
        return int(mask.shape[0])
    raise ValueError(f"unsupported MASK tensor shape {tuple(mask.shape)}")


def _validate_optional_mask_batch(name: str, value: Optional[torch.Tensor], batch_size: int) -> None:
    if value is None:
        return
    actual = _mask_batch_size(value)
    if actual not in {1, batch_size, max(1, batch_size - 1)}:
        raise ValueError(
            f"{name} batch size must be 1, the full frame count, or frame_count - 1; "
            f"got {actual} for {batch_size} frames"
        )


def _extract_optional_mask(
    name: str,
    value: Optional[torch.Tensor],
    frame_index: int,
    batch_size: int,
) -> Optional[torch.Tensor]:
    if value is None:
        return None

    if value.ndim == 2:
        return value.unsqueeze(0)

    if value.ndim == 3:
        actual = int(value.shape[0])
        if actual == 1:
            return value
        if actual == batch_size:
            return value[frame_index : frame_index + 1]
        if actual == batch_size - 1:
            if frame_index == 0:
                return None
            return value[frame_index - 1 : frame_index]

    if value.ndim == 4 and value.shape[-1] == 1:
        actual = int(value.shape[0])
        if actual == 1:
            return value[..., 0]
        if actual == batch_size:
            return value[frame_index : frame_index + 1, ..., 0]
        if actual == batch_size - 1:
            if frame_index == 0:
                return None
            return value[frame_index - 1 : frame_index, ..., 0]

    raise ValueError(f"{name} batch size mismatch")


def _mask_to_image_weight(mask: Optional[torch.Tensor], reference_image: torch.Tensor) -> Optional[torch.Tensor]:
    if mask is None:
        return None

    if mask.ndim == 2:
        mask = mask.unsqueeze(0)
    if mask.ndim == 4 and mask.shape[-1] == 1:
        mask = mask[..., 0]
    if mask.ndim != 3:
        raise ValueError(f"expected mask with shape [B,H,W], got {tuple(mask.shape)}")

    mask = mask.to(device=reference_image.device, dtype=reference_image.dtype)
    if mask.shape[0] == 1 and reference_image.shape[0] != 1:
        mask = mask.expand(reference_image.shape[0], -1, -1)
    if mask.shape[1:3] != reference_image.shape[1:3]:
        mask = F.interpolate(
            mask.unsqueeze(1),
            size=reference_image.shape[1:3],
            mode="bilinear",
            align_corners=False,
        ).squeeze(1)

    return mask.unsqueeze(-1).clamp(0.0, 1.0)


def _ensure_same_image_shape(current: torch.Tensor, previous: torch.Tensor, context: str) -> torch.Tensor:
    if previous.shape == current.shape:
        return previous.to(device=current.device, dtype=current.dtype)

    if previous.shape[0] != current.shape[0] or previous.shape[-1] != current.shape[-1]:
        raise ValueError(
            f"{context} has incompatible IMAGE shape {tuple(previous.shape)} "
            f"for current frame {tuple(current.shape)}"
        )

    resized = previous.permute(0, 3, 1, 2)
    resized = F.interpolate(resized, size=current.shape[1:3], mode="bilinear", align_corners=False)
    return resized.permute(0, 2, 3, 1).to(device=current.device, dtype=current.dtype)


def _frame_difference_score(previous_frame: torch.Tensor, current_frame: torch.Tensor) -> float:
    aligned_previous = _ensure_same_image_shape(current_frame, previous_frame, "previous input")
    return float(torch.mean(torch.abs(current_frame - aligned_previous)).item())


def _compute_reset(
    frame_index: int,
    previous_input: Optional[torch.Tensor],
    current_input: torch.Tensor,
    reset_every_n: int,
    scene_cut_threshold: float,
) -> bool:
    if frame_index == 0 or previous_input is None:
        return True
    if reset_every_n > 0 and frame_index % reset_every_n == 0:
        return True
    if scene_cut_threshold > 0.0 and _frame_difference_score(previous_input, current_input) >= scene_cut_threshold:
        return True
    return False


def _blend_images(
    current: torch.Tensor,
    prior: torch.Tensor,
    strength: float,
    confidence: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    prior = _ensure_same_image_shape(current, prior, "temporal prior")
    alpha = torch.full_like(current[..., :1], float(strength))
    if confidence is not None:
        alpha = alpha * _mask_to_image_weight(confidence, current)
    alpha = alpha.clamp(0.0, 1.0)
    return current * (1.0 - alpha) + prior * alpha


def _cv2_require() -> Any:
    if cv2 is None:
        raise RuntimeError("internal_flow_warp requires OpenCV (cv2) in the ComfyUI Python environment")
    return cv2


def _image_to_numpy_uint8(frame: torch.Tensor):
    return frame.detach().cpu().clamp(0.0, 1.0).mul(255.0).round().to(torch.uint8).numpy()


def _compute_backward_flow(current_frame: torch.Tensor, previous_frame: torch.Tensor):
    cv2_mod = _cv2_require()
    current_np = _image_to_numpy_uint8(current_frame[0])
    previous_np = _image_to_numpy_uint8(previous_frame[0])
    current_gray = cv2_mod.cvtColor(current_np, cv2_mod.COLOR_RGB2GRAY)
    previous_gray = cv2_mod.cvtColor(previous_np, cv2_mod.COLOR_RGB2GRAY)
    return cv2_mod.calcOpticalFlowFarneback(
        current_gray,
        previous_gray,
        None,
        pyr_scale=0.5,
        levels=3,
        winsize=15,
        iterations=3,
        poly_n=5,
        poly_sigma=1.2,
        flags=0,
    )


def _warp_with_flow(image: torch.Tensor, backward_flow) -> torch.Tensor:
    cv2_mod = _cv2_require()
    import numpy as np

    image_np = image[0].detach().cpu().numpy().astype("float32", copy=False)
    height, width = backward_flow.shape[:2]
    grid_x, grid_y = np.meshgrid(range(width), range(height))
    map_x = (grid_x.astype("float32") + backward_flow[..., 0]).astype("float32")
    map_y = (grid_y.astype("float32") + backward_flow[..., 1]).astype("float32")
    warped = cv2_mod.remap(
        image_np,
        map_x,
        map_y,
        interpolation=cv2_mod.INTER_LINEAR,
        borderMode=cv2_mod.BORDER_REPLICATE,
    )
    return torch.from_numpy(warped).unsqueeze(0).to(device=image.device, dtype=image.dtype)


def _photometric_flow_confidence(
    current_frame: torch.Tensor,
    previous_input_frame: torch.Tensor,
    backward_flow,
) -> torch.Tensor:
    warped_previous_input = _warp_with_flow(previous_input_frame, backward_flow)
    error = torch.mean(torch.abs(current_frame - warped_previous_input), dim=-1, keepdim=True)
    return (1.0 - error).clamp(0.0, 1.0)


def _resolve_temporal_prior(
    temporal_mode: str,
    frame_index: int,
    current_frame: torch.Tensor,
    previous_input: Optional[torch.Tensor],
    previous_output: Optional[torch.Tensor],
    warped_previous_images: Optional[torch.Tensor],
    flow_confidence_mask: Optional[torch.Tensor],
    flow_confidence_scale: float,
    temporal_strength: float,
    batch_size: int,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
    if temporal_mode == "off" or frame_index == 0 or previous_output is None:
        return current_frame, None

    if temporal_mode == "prev_output_blend":
        return _blend_images(current_frame, previous_output, temporal_strength), None

    external_mask = None
    if temporal_mode in {"external_warped_prev", "internal_flow_warp"}:
        external_mask = _extract_optional_mask("flow_confidence", flow_confidence_mask, frame_index, batch_size)
        if external_mask is not None:
            external_mask = external_mask * float(flow_confidence_scale)
        elif flow_confidence_scale < 1.0:
            external_mask = torch.full(
                (1, current_frame.shape[1], current_frame.shape[2]),
                float(flow_confidence_scale),
                device=current_frame.device,
                dtype=current_frame.dtype,
            )

    if temporal_mode == "external_warped_prev":
        warped_frame = _extract_optional_frame(
            "warped_previous_images",
            warped_previous_images,
            frame_index,
            batch_size,
        )
        if warped_frame is None:
            raise ValueError("external_warped_prev requires warped_previous_images after the first frame")
        warped_frame = _temporal_rgb_view(warped_frame)
        return _blend_images(current_frame, warped_frame, temporal_strength, confidence=external_mask), warped_frame

    if temporal_mode == "internal_flow_warp":
        if previous_input is None:
            return current_frame, None
        previous_input = _ensure_same_image_shape(current_frame, previous_input, "previous input")
        previous_output = _ensure_same_image_shape(current_frame, previous_output, "previous output")
        backward_flow = _compute_backward_flow(current_frame, previous_input)
        warped_previous_output = _warp_with_flow(previous_output, backward_flow)
        confidence = _photometric_flow_confidence(current_frame, previous_input, backward_flow)
        if external_mask is not None:
            confidence = confidence * _mask_to_image_weight(external_mask, current_frame)
        return (
            _blend_images(current_frame, warped_previous_output, temporal_strength, confidence=confidence),
            warped_previous_output,
        )

    raise ValueError(f"unsupported temporal_mode: {temporal_mode}")


def _seed_for_frame(noise: Any, frame_index: int, lock_seed: bool, seed_stride_mode: str) -> Optional[int]:
    if not lock_seed:
        return None
    if not hasattr(noise, "seed"):
        raise TypeError("lock_seed=True requires a NOISE input with a 'seed' attribute")

    base_seed = int(noise.seed)
    if seed_stride_mode == "fixed":
        return base_seed
    if seed_stride_mode == "sequential":
        return base_seed + frame_index
    raise ValueError(f"unsupported seed_stride_mode: {seed_stride_mode}")


def _batch_index_for_frame(frame_index: int, lock_seed: bool) -> int:
    if lock_seed:
        return 0
    return int(frame_index)


class FluxVideoCleanupTemporalAdvanced:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "vae": ("VAE",),
                "noise": ("NOISE",),
                "guider": ("GUIDER",),
                "sampler": ("SAMPLER",),
                "sigmas": ("SIGMAS",),
                "temporal_mode": (("off", "prev_output_blend", "external_warped_prev", "internal_flow_warp"),),
                "temporal_strength": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01}),
                "reset_every_n": ("INT", {"default": 0, "min": 0, "max": 4096}),
                "scene_cut_threshold": ("FLOAT", {"default": 0.18, "min": 0.0, "max": 1.0, "step": 0.01}),
                "flow_confidence_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "advanced": True}),
                "lock_seed": ("BOOLEAN", {"default": True}),
                "seed_stride_mode": (("fixed", "sequential"),),
                "encode_chunk_size": ("INT", {"default": 8, "min": 1, "max": 4096}),
                "sample_chunk_size": (
                    "INT",
                    {
                        "default": 1,
                        "min": 1,
                        "max": 4096,
                        "tooltip": "Only used when temporal_mode is off. Temporal recurrence forces per-frame sampling.",
                    },
                ),
                "decode_chunk_size": ("INT", {"default": 8, "min": 1, "max": 4096}),
                "auto_reduce_on_oom": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": "Only used when temporal_mode is off. Temporal recurrence always samples one frame at a time.",
                    },
                ),
                "min_chunk_size": (
                    "INT",
                    {
                        "default": 1,
                        "min": 1,
                        "max": 4096,
                        "tooltip": "Only used when temporal_mode is off. Temporal recurrence always samples one frame at a time.",
                    },
                ),
                "clear_cache_on_retry": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "advanced": True,
                        "tooltip": "Only used when temporal_mode is off. Temporal recurrence does not use chunk-size retry fallback.",
                    },
                ),
                "clear_cache_between_chunks": ("BOOLEAN", {"default": False, "advanced": True}),
            },
            "optional": {
                "warped_previous_images": ("IMAGE",),
                "flow_confidence": ("MASK",),
            },
        }

    RETURN_TYPES = ("IMAGE", "LATENT", "LATENT", "MASK")
    RETURN_NAMES = ("images", "latents", "denoised_latents", "reset_mask")
    FUNCTION = "cleanup"
    CATEGORY = "sampling/custom_sampling"

    def cleanup(
        self,
        images,
        vae,
        noise,
        guider,
        sampler,
        sigmas,
        temporal_mode,
        temporal_strength,
        reset_every_n,
        scene_cut_threshold,
        flow_confidence_scale,
        lock_seed,
        seed_stride_mode,
        encode_chunk_size,
        sample_chunk_size,
        decode_chunk_size,
        auto_reduce_on_oom=True,
        min_chunk_size=1,
        clear_cache_on_retry=True,
        clear_cache_between_chunks=False,
        warped_previous_images=None,
        flow_confidence=None,
    ):
        batch_size = _frame_batch_size(images)
        if batch_size == 0:
            raise ValueError("images batch must not be empty")

        if temporal_mode == "off":
            latents = encode_image_batch_chunked(
                images,
                vae,
                encode_chunk_size,
                clear_cache_between_chunks=clear_cache_between_chunks,
            )
            sampled, denoised = sample_latent_batch_chunked(
                noise,
                guider,
                sampler,
                sigmas,
                latents,
                config=SampleChunkConfig(
                    chunk_size=sample_chunk_size,
                    auto_reduce_on_oom=auto_reduce_on_oom,
                    min_chunk_size=min_chunk_size,
                    clear_cache_on_retry=clear_cache_on_retry,
                    clear_cache_between_chunks=clear_cache_between_chunks,
                ),
            )
            decoded = decode_latent_batch_chunked(
                sampled,
                vae,
                decode_chunk_size,
                clear_cache_between_chunks=clear_cache_between_chunks,
            )
            reset_mask = torch.ones(
                (batch_size, images.shape[1], images.shape[2]),
                device=images.device,
                dtype=images.dtype,
            )
            return decoded, sampled, denoised, reset_mask

        if temporal_mode == "external_warped_prev":
            _validate_optional_frame_batch("warped_previous_images", warped_previous_images, batch_size)
        if temporal_mode in {"external_warped_prev", "internal_flow_warp"}:
            _validate_optional_mask_batch("flow_confidence", flow_confidence, batch_size)

        previous_input = None
        previous_output = None
        output_images = []
        output_latents = []
        output_denoised_latents = []
        reset_values = []

        for frame_index in range(batch_size):
            current_frame = _temporal_rgb_view(_single_frame(images, frame_index))
            should_reset = _compute_reset(
                frame_index,
                previous_input,
                current_frame,
                int(reset_every_n),
                float(scene_cut_threshold),
            )
            reset_values.append(float(should_reset))

            if should_reset:
                prior_frame = current_frame
            else:
                prior_frame, _ = _resolve_temporal_prior(
                    temporal_mode,
                    frame_index,
                    current_frame,
                    previous_input,
                    previous_output,
                    warped_previous_images,
                    flow_confidence,
                    float(flow_confidence_scale),
                    float(temporal_strength),
                    batch_size,
                )

            encoded = encode_image_batch_chunked(
                prior_frame,
                vae,
                encode_chunk_size,
                clear_cache_between_chunks=False,
            )
            sampled, denoised = sample_single_latent_temporal(
                noise,
                guider,
                sampler,
                sigmas,
                encoded,
                noise_seed_override=_seed_for_frame(noise, frame_index, bool(lock_seed), seed_stride_mode),
                forced_batch_index_value=_batch_index_for_frame(frame_index, bool(lock_seed)),
                clear_cache_between_chunks=clear_cache_between_chunks,
            )
            decoded = decode_latent_batch_chunked(
                sampled,
                vae,
                decode_chunk_size,
                clear_cache_between_chunks=False,
            )

            output_images.append(decoded)
            output_latents.append(sampled["samples"])
            output_denoised_latents.append(denoised["samples"])

            previous_input = current_frame
            previous_output = decoded

            if clear_cache_between_chunks and frame_index + 1 < batch_size:
                soft_clear_cache()

        images_out = torch.cat(output_images, dim=0)
        latents_out = {"samples": torch.cat(output_latents, dim=0)}
        denoised_out = {"samples": torch.cat(output_denoised_latents, dim=0)}
        reset_mask = torch.tensor(reset_values, device=images.device, dtype=images.dtype).view(batch_size, 1, 1)
        reset_mask = reset_mask.expand(batch_size, images.shape[1], images.shape[2])
        return images_out, latents_out, denoised_out, reset_mask


NODE_CLASS_MAPPINGS = {
    "FluxVideoCleanupTemporalAdvanced": FluxVideoCleanupTemporalAdvanced,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "FluxVideoCleanupTemporalAdvanced": "Flux Video Cleanup Temporal Advanced",
}