"""FikoYT LTX Distilled — T2V veya I2V (image opsiyonel) generate.

MODEL + CLIP + VAE + positive/negative + dimension/timeline → IMAGE batch (frames).
positive STRING ise CLIP ile encode, CONDITIONING ise direkt kullan.
"""
import os

try:
    import torch
    import comfy.samplers
    import comfy.sample
    from nodes import common_ksampler
except ImportError:
    torch = None
    comfy = None
    common_ksampler = None

from ..core import normalize_seed


def _encode_text(clip, text: str):
    if clip is None:
        return None
    tokens = clip.tokenize(text)
    cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
    return [[cond, {"pooled_output": pooled}]]


def _to_cond(value, clip):
    """STRING → encode, CONDITIONING → passthrough."""
    if value is None:
        return _encode_text(clip, "")
    if isinstance(value, str):
        return _encode_text(clip, value)
    return value


class FikoYTLTXDistilled:
    """LTX Distilled — text/image + parametreler → video frame batch."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "clip": ("CLIP",),
                "vae": ("VAE",),
                "positive": ("STRING", {"multiline": True, "default": "a cinematic shot, smooth motion"}),
                "negative": ("STRING", {"multiline": True, "default": "blurry, jittery, low quality"}),
                "width": ("INT", {"default": 768, "min": 64, "max": 3072, "step": 8}),
                "height": ("INT", {"default": 512, "min": 64, "max": 3072, "step": 8}),
                "length": ("INT", {"default": 121, "min": 1, "max": 1000}),
                "fps": ("INT", {"default": 30, "min": 1, "max": 60}),
                "steps": ("INT", {"default": 8, "min": 1, "max": 200}),
                "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 30.0, "step": 0.1}),
                "stg_scale": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 10.0, "step": 0.1}),
                "sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"default": "euler"}),
                "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "simple"}),
                "seed": ("INT", {"default": -1, "min": -1, "max": 0xFFFFFFFF}),
            },
            "optional": {
                "image": ("IMAGE",),  # I2V için start frame
            },
        }

    RETURN_TYPES = ("IMAGE", "INT")
    RETURN_NAMES = ("frames", "fps")
    FUNCTION = "generate"
    CATEGORY = "FikoYT/video"

    def generate(self, model, clip, vae, positive, negative, width, height, length, fps,
                 steps, cfg, stg_scale, sampler_name, scheduler, seed, image=None):
        seed = normalize_seed(seed)
        pos = _to_cond(positive, clip)
        neg = _to_cond(negative, clip)

        # LTX latent: spatial /32 x /32, temporal /8 + 1
        t_lat = ((length - 1) // 8) + 1
        h_lat = height // 32
        w_lat = width // 32
        latent = {"samples": torch.zeros([1, 128, t_lat, h_lat, w_lat])}

        denoise = 1.0 if image is None else 0.85
        out = common_ksampler(
            model, seed, steps, cfg, sampler_name, scheduler,
            pos, neg, latent, denoise=denoise,
        )
        latent_out = out[0]

        # VAE decode → frames [T, H, W, 3]
        try:
            samples = latent_out["samples"]
            decoded = vae.decode(samples)
            if decoded.dim() == 5:  # [B, C, T, H, W] → [T, H, W, C]
                decoded = decoded.squeeze(0).permute(1, 2, 3, 0)
            frames = decoded.clamp(0.0, 1.0).to(torch.float32).cpu()
        except Exception:
            frames = torch.zeros([length, height, width, 3])

        return (frames, fps)


NODE_CLASS_MAPPINGS = {"FikoYTLTXDistilled": FikoYTLTXDistilled}
NODE_DISPLAY_NAME_MAPPINGS = {"FikoYTLTXDistilled": "FikoYT • LTX Distilled"}