import torch
import comfy.utils
import comfy_extras.nodes_lt as nodes_lt
import comfy_extras.nodes_post_processing as post_processing


class LTXVImgToVideoInplaceAtIndex:
    """
    Places an encoded image at any frame position in an existing video latent.

    Improvements over the built-in LTXVImgToVideoInplace:
    - Arbitrary frame_idx (not just frame 0). Negative values count from the end.
    - Noise mask chaining: uses get_noise_mask() so multiple instances can be
      chained without overwriting each other's masks.
    - Preprocessing options (CRF, blur, interpolation, crop) matching
      LTXVAddGuideAdvanced.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "vae": ("VAE",),
                "image": ("IMAGE",),
                "latent": ("LATENT",),
                "frame_idx": (
                    "INT",
                    {
                        "default": 0,
                        "min": -9999,
                        "max": 9999,
                        "step": 1,
                        "tooltip": (
                            "Pixel-space frame index. 0 = first frame, -1 = last frame. "
                            "Converted to latent index using VAE temporal downscale factor."
                        ),
                    },
                ),
                "strength": (
                    "FLOAT",
                    {
                        "default": 1.0,
                        "min": 0.0,
                        "max": 1.0,
                        "step": 0.01,
                        "tooltip": "Conditioning strength. 1.0 = fully conditioned, 0.0 = no conditioning.",
                    },
                ),
                "crf": (
                    "INT",
                    {
                        "default": 29,
                        "min": 0,
                        "max": 51,
                        "step": 1,
                        "tooltip": "CRF compression. 0 = disabled. Higher values add more compression artifacts.",
                    },
                ),
                "blur_radius": (
                    "INT",
                    {
                        "default": 0,
                        "min": 0,
                        "max": 7,
                        "step": 1,
                        "tooltip": "Gaussian blur kernel radius. 0 = disabled.",
                    },
                ),
                "interpolation": (
                    [
                        "lanczos",
                        "bislerp",
                        "nearest",
                        "bilinear",
                        "bicubic",
                        "area",
                        "nearest-exact",
                    ],
                    {"default": "lanczos"},
                ),
                "crop": (["disabled", "center"], {"default": "disabled"}),
                "bypass": (
                    "BOOLEAN",
                    {"default": False, "tooltip": "Bypass the conditioning."},
                ),
            }
        }

    RETURN_TYPES = ("LATENT",)
    RETURN_NAMES = ("latent",)
    FUNCTION = "generate"
    CATEGORY = "conditioning/video_models"
    DESCRIPTION = (
        "Places an image at an arbitrary frame position in a video latent. "
        "Supports chaining: noise masks from previous calls are preserved."
    )

    def generate(
        self, vae, image, latent, frame_idx, strength, crf, blur_radius,
        interpolation, crop, bypass=False,
    ):
        if bypass:
            return (latent,)

        # VAE scale factors
        time_scale, height_scale, width_scale = vae.downscale_index_formula

        # Compute target pixel dimensions from latent shape
        samples = latent["samples"]
        batch, channels, latent_frames, latent_h, latent_w = samples.shape
        width = latent_w * width_scale
        height = latent_h * height_scale

        # Preprocessing: resize, CRF, blur (matching LTXVAddGuideAdvanced)
        image = (
            comfy.utils.common_upscale(
                image.movedim(-1, 1), width, height, interpolation, crop=crop
            )
            .movedim(1, -1)
            .clamp(0, 1)
        )
        if crf > 0:
            image = nodes_lt.LTXVPreprocess.execute(image, crf)[0]
        if blur_radius > 0:
            sigma = 0.3 * blur_radius
            image = post_processing.Blur.execute(image, blur_radius, sigma)[0]

        # VAE encode (RGB channels only)
        encode_pixels = image[:, :, :, :3]
        t = vae.encode(encode_pixels)
        cond_length = t.shape[2]

        # Convert pixel frame_idx to latent index
        total_pixel_frames = (latent_frames - 1) * time_scale + 1

        if frame_idx < 0:
            frame_idx = max(0, total_pixel_frames + frame_idx)
        frame_idx = min(frame_idx, total_pixel_frames - 1)

        if frame_idx == 0:
            latent_idx = 0
        else:
            latent_idx = ((frame_idx - 1) // time_scale) + 1

        # Validate bounds
        assert latent_idx + cond_length <= latent_frames, (
            f"Conditioning frames exceed latent length. "
            f"latent_idx={latent_idx}, cond_length={cond_length}, latent_frames={latent_frames}"
        )

        # Chain noise mask properly (preserves masks from previous calls)
        noise_mask = nodes_lt.get_noise_mask(latent)
        samples = samples.clone()

        # Replace at the target index
        samples[:, :, latent_idx:latent_idx + cond_length] = t
        noise_mask[:, :, latent_idx:latent_idx + cond_length] = 1.0 - strength

        return ({"samples": samples, "noise_mask": noise_mask},)