"""FikoYT Video Grid — 2-4 IMAGE batch'i grid (2x2 / 3x1 / 1x3 / 4x1) olarak birleştir.

Her frame'de tile layout. Frame sayısı en kısa batch'e göre kırpılır.
"""
try:
    import torch
    import torch.nn.functional as F
except ImportError:
    torch = None
    F = None


def _resize_to(x, h, w):
    if x.shape[1] == h and x.shape[2] == w:
        return x
    y = x.permute(0, 3, 1, 2)
    y = F.interpolate(y, size=(h, w), mode="bilinear", align_corners=False)
    return y.permute(0, 2, 3, 1).clamp(0.0, 1.0)


def _trim_to(batches, n):
    return [b[:n] if b is not None else None for b in batches]


def _layout(batches, grid):
    """batches: list of IMAGE [T,H,W,C], grid: '2x2'|'3x1'|'1x3'|'4x1'."""
    valid = [b for b in batches if b is not None]
    if not valid:
        return torch.zeros([1, 64, 64, 3])

    n = min(b.shape[0] for b in valid)
    valid = _trim_to(valid, n)

    h = valid[0].shape[1]
    w = valid[0].shape[2]
    # Hepsini ilk batch boyutuna eşitle
    valid = [_resize_to(b, h, w) for b in valid]

    if grid == "2x2":
        while len(valid) < 4:
            valid.append(torch.zeros_like(valid[0]))
        top = torch.cat([valid[0], valid[1]], dim=2)
        bot = torch.cat([valid[2], valid[3]], dim=2)
        out = torch.cat([top, bot], dim=1)
    elif grid == "3x1":  # 3 yan yana, tek satır
        while len(valid) < 3:
            valid.append(torch.zeros_like(valid[0]))
        out = torch.cat(valid[:3], dim=2)
    elif grid == "1x3":  # 3 alt alta, tek sütun
        while len(valid) < 3:
            valid.append(torch.zeros_like(valid[0]))
        out = torch.cat(valid[:3], dim=1)
    elif grid == "4x1":  # 4 yan yana
        while len(valid) < 4:
            valid.append(torch.zeros_like(valid[0]))
        out = torch.cat(valid[:4], dim=2)
    else:
        out = valid[0]
    return out


class FikoYTVideoGrid:
    """1-4 IMAGE batch → grid layout IMAGE batch."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images_1": ("IMAGE",),
                "grid": (["2x2", "3x1", "1x3", "4x1"], {"default": "2x2"}),
            },
            "optional": {
                "images_2": ("IMAGE",),
                "images_3": ("IMAGE",),
                "images_4": ("IMAGE",),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("frames",)
    FUNCTION = "build"
    CATEGORY = "FikoYT/video"

    def build(self, images_1, grid, images_2=None, images_3=None, images_4=None):
        out = _layout([images_1, images_2, images_3, images_4], grid)
        return (out,)


NODE_CLASS_MAPPINGS = {"FikoYTVideoGrid": FikoYTVideoGrid}
NODE_DISPLAY_NAME_MAPPINGS = {"FikoYTVideoGrid": "FikoYT • Video Grid"}