import sys
import types
from pathlib import Path

import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))


class FakeNestedTensor:
    def __init__(self, tensors):
        self.tensors = list(tensors)
        self.is_nested = True

    def unbind(self):
        return self.tensors

    @property
    def shape(self):
        return self.tensors[0].shape

    @property
    def device(self):
        return self.tensors[0].device

    @property
    def dtype(self):
        return self.tensors[0].dtype

    @property
    def layout(self):
        return self.tensors[0].layout

    def size(self):
        return self.tensors[0].size()


sample_calls = []


def _zeros_like(value):
    if getattr(value, "is_nested", False):
        return FakeNestedTensor([torch.zeros_like(t) for t in value.unbind()])
    return torch.zeros_like(value)


def _sample_like(value, mask, seed):
    if getattr(value, "is_nested", False):
        masks = mask.unbind() if getattr(mask, "is_nested", False) else [None] * len(value.unbind())
        return FakeNestedTensor([
            _sample_like(tensor, masks[index], seed)
            for index, tensor in enumerate(value.unbind())
        ])
    if mask is None:
        return value + float(seed)
    return value + mask.to(value.device, value.dtype) * float(seed)


def fake_prepare_noise(samples, seed, noise_inds=None):
    return _zeros_like(samples)


def fake_sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image,
                       noise_mask=None, callback=None, disable_pbar=False, seed=None):
    sample_calls.append({
        "seed": seed,
        "positive": positive,
        "negative": negative,
        "cfg": cfg,
        "latent_image": _sample_snapshot(latent_image),
        "noise_mask": _sample_snapshot(noise_mask),
    })
    return _sample_like(latent_image, noise_mask, seed)


def _sample_snapshot(value):
    if value is None:
        return None
    if getattr(value, "is_nested", False):
        return FakeNestedTensor([tensor.detach().clone() for tensor in value.unbind()])
    return value.detach().clone()


comfy_mod = types.ModuleType("comfy")
comfy_kd_mod = types.ModuleType("comfy.k_diffusion")
comfy_kd_sampling_mod = types.ModuleType("comfy.k_diffusion.sampling")
comfy_model_sampling_mod = types.ModuleType("comfy.model_sampling")
comfy_nested_mod = types.ModuleType("comfy.nested_tensor")
comfy_sample_mod = types.ModuleType("comfy.sample")
comfy_samplers_mod = types.ModuleType("comfy.samplers")
comfy_utils_mod = types.ModuleType("comfy.utils")

comfy_nested_mod.NestedTensor = FakeNestedTensor
comfy_sample_mod.prepare_noise = fake_prepare_noise
comfy_sample_mod.sample_custom = fake_sample_custom

comfy_mod.k_diffusion = comfy_kd_mod
comfy_kd_mod.sampling = comfy_kd_sampling_mod
comfy_mod.model_sampling = comfy_model_sampling_mod
comfy_mod.nested_tensor = comfy_nested_mod
comfy_mod.sample = comfy_sample_mod
comfy_mod.samplers = comfy_samplers_mod
comfy_mod.utils = comfy_utils_mod

sys.modules.setdefault("comfy", comfy_mod)
sys.modules.setdefault("comfy.k_diffusion", comfy_kd_mod)
sys.modules.setdefault("comfy.k_diffusion.sampling", comfy_kd_sampling_mod)
sys.modules.setdefault("comfy.model_sampling", comfy_model_sampling_mod)
sys.modules.setdefault("comfy.nested_tensor", comfy_nested_mod)
sys.modules.setdefault("comfy.sample", comfy_sample_mod)
sys.modules.setdefault("comfy.samplers", comfy_samplers_mod)
sys.modules.setdefault("comfy.utils", comfy_utils_mod)

from samplers import FunPackLTXAVSceneChainSampler


class FakeVAE:
    downscale_index_formula = (1, 1, 1)


def scene_cond(index):
    return (
        torch.ones(1, 2, 3) * float(index + 1),
        {"funpack_scene_text": f"scene {index + 1}"},
    )


def test_scene_chain_detects_scene_count_and_increments_seed():
    sample_calls.clear()
    node = FunPackLTXAVSceneChainSampler()
    latent_template = {"samples": torch.zeros(1, 2, 5, 3, 3)}
    positive = [scene_cond(0), scene_cond(1), scene_cond(2)]
    negative = [(torch.zeros(1, 2, 3), {})]

    latent, status, scene_count, report = node.sample(
        model=object(),
        vae=FakeVAE(),
        positive=positive,
        negative=negative,
        sampler=object(),
        sigmas=torch.tensor([1.0, 0.0]),
        seed=10,
        latent_template=latent_template,
        num_frames_per_scene=5,
        frame_overlap=2,
        cfg=1.5,
        max_scenes=8,
    )

    assert scene_count == 3
    assert [call["seed"] for call in sample_calls] == [10, 11, 12]
    assert [call["positive"][0][1]["funpack_scene_text"] for call in sample_calls] == ["scene 1", "scene 2", "scene 3"]
    assert latent["samples"].shape[2] == 11
    assert "Scene chain complete" in status
    assert "Scene 3" in report


def test_scene_chain_accepts_manual_combined_conditioning_without_metadata():
    sample_calls.clear()
    node = FunPackLTXAVSceneChainSampler()
    latent_template = {"samples": torch.zeros(1, 2, 5, 3, 3)}
    positive = [
        (torch.ones(1, 2, 3), {}),
        (torch.ones(1, 2, 3) * 2.0, {}),
        (torch.ones(1, 2, 3) * 3.0, {}),
    ]

    _, status, scene_count, report = node.sample(
        model=object(),
        vae=FakeVAE(),
        positive=positive,
        negative=[],
        sampler=object(),
        sigmas=torch.tensor([1.0, 0.0]),
        seed=70,
        latent_template=latent_template,
        num_frames_per_scene=5,
        frame_overlap=2,
        cfg=1.0,
        max_scenes=8,
    )

    assert scene_count == 3
    assert [call["seed"] for call in sample_calls] == [70, 71, 72]
    assert [call["positive"][0][0].mean().item() for call in sample_calls] == [1.0, 2.0, 3.0]
    assert "3 scene(s)" in status
    assert "Scene 1: seed=70, text=Scene 1" in report
    assert "Scene 3: seed=72, text=Scene 3" in report


def test_scene_chain_uses_scene_seed_metadata():
    sample_calls.clear()
    node = FunPackLTXAVSceneChainSampler()
    latent_template = {"samples": torch.zeros(1, 2, 5, 3, 3)}
    positive = [
        (torch.ones(1, 2, 3), {"funpack_scene_text": "scene a", "funpack_scene_seed": 101}),
        (torch.ones(1, 2, 3), {"funpack_scene_text": "scene b", "funpack_scene_seed": 202}),
        (torch.ones(1, 2, 3), {"funpack_scene_text": "scene c", "funpack_scene_seed": 303}),
    ]

    _, _, scene_count, report = node.sample(
        model=object(),
        vae=FakeVAE(),
        positive=positive,
        negative=[],
        sampler=object(),
        sigmas=torch.tensor([1.0, 0.0]),
        seed=10,
        latent_template=latent_template,
        num_frames_per_scene=5,
        frame_overlap=2,
        cfg=1.0,
        max_scenes=8,
    )

    assert scene_count == 3
    assert [call["seed"] for call in sample_calls] == [101, 202, 303]
    assert "seed=202" in report


def test_scene_chain_use_same_seed_forces_first_seed():
    sample_calls.clear()
    node = FunPackLTXAVSceneChainSampler()
    latent_template = {"samples": torch.zeros(1, 2, 5, 3, 3)}
    positive = [
        (torch.ones(1, 2, 3), {"funpack_scene_text": "scene a", "funpack_scene_seed": 101}),
        (torch.ones(1, 2, 3), {"funpack_scene_text": "scene b", "funpack_scene_seed": 202}),
        (torch.ones(1, 2, 3), {"funpack_scene_text": "scene c"}),
    ]

    _, _, scene_count, _ = node.sample(
        model=object(),
        vae=FakeVAE(),
        positive=positive,
        negative=[],
        sampler=object(),
        sigmas=torch.tensor([1.0, 0.0]),
        seed=10,
        latent_template=latent_template,
        num_frames_per_scene=5,
        frame_overlap=2,
        cfg=1.0,
        max_scenes=8,
        use_same_seed=True,
    )

    assert scene_count == 3
    assert [call["seed"] for call in sample_calls] == [101, 101, 101]


def test_scene_chain_preserves_nested_av_structure_and_audio_length():
    sample_calls.clear()
    node = FunPackLTXAVSceneChainSampler()
    video = torch.zeros(1, 2, 5, 3, 3)
    audio = torch.zeros(1, 1, 10, 4)
    latent_template = {"samples": FakeNestedTensor([video, audio])}
    positive = [scene_cond(0), scene_cond(1)]

    latent, _, scene_count, _ = node.sample(
        model=object(),
        vae=FakeVAE(),
        positive=positive,
        negative=[],
        sampler=object(),
        sigmas=torch.tensor([1.0, 0.0]),
        seed=20,
        latent_template=latent_template,
        num_frames_per_scene=5,
        frame_overlap=2,
        cfg=1.0,
        max_scenes=8,
    )

    video_out, audio_out = latent["samples"].unbind()
    assert scene_count == 2
    assert video_out.shape[2] == 8
    assert audio_out.shape[2] == 16


def test_scene_chain_default_max_is_eight_but_allows_more():
    inputs = FunPackLTXAVSceneChainSampler.INPUT_TYPES()["required"]["max_scenes"][1]
    assert inputs["default"] == 8
    assert "max" not in inputs

    sample_calls.clear()
    node = FunPackLTXAVSceneChainSampler()
    latent_template = {"samples": torch.zeros(1, 2, 3, 2, 2)}
    positive = [scene_cond(index) for index in range(10)]

    latent, status, scene_count, report = node.sample(
        model=object(),
        vae=FakeVAE(),
        positive=positive,
        negative=[],
        sampler=object(),
        sigmas=torch.tensor([1.0, 0.0]),
        seed=30,
        latent_template=latent_template,
        num_frames_per_scene=3,
        frame_overlap=0,
        cfg=1.0,
        max_scenes=10,
    )

    assert scene_count == 10
    assert len(sample_calls) == 10
    assert sample_calls[-1]["seed"] == 39
    assert latent["samples"].shape[2] == 30
    assert "10 scene(s)" in status
    assert "Scene 10" in report


def test_scene_chain_can_append_i2v_template_as_hidden_guide():
    sample_calls.clear()
    node = FunPackLTXAVSceneChainSampler()
    samples = torch.zeros(1, 2, 5, 1, 1)
    samples[:, :, 0] = 7.0
    mask = torch.ones(1, 1, 5, 1, 1)
    mask[:, :, 0] = 0.0
    latent_template = {"samples": samples, "noise_mask": mask}
    positive = [scene_cond(0), scene_cond(1)]

    latent, status, scene_count, _ = node.sample(
        model=object(),
        vae=FakeVAE(),
        positive=positive,
        negative=[],
        sampler=object(),
        sigmas=torch.tensor([1.0, 0.0]),
        seed=40,
        latent_template=latent_template,
        num_frames_per_scene=5,
        frame_overlap=2,
        cfg=1.0,
        max_scenes=2,
        carry_i2v_guides=True,
    )

    second_call = sample_calls[1]
    assert scene_count == 2
    assert second_call["latent_image"].shape[2] == 6
    assert torch.all(second_call["latent_image"][:, :, 2:5] == 0.0)
    assert torch.all(second_call["latent_image"][:, :, 5] == 7.0)
    assert torch.all(second_call["noise_mask"][:, :, :2] == 0.0)
    assert torch.all(second_call["noise_mask"][:, :, 2:5] == 1.0)
    assert torch.all(second_call["noise_mask"][:, :, 5] == 0.0)
    assert "keyframe_idxs" in second_call["positive"][0][1]
    assert second_call["positive"][0][1]["guide_attention_entries"][0]["pre_filter_count"] == 1
    assert latent["samples"].shape[2] == 8
    assert "i2v guide tokens=1 latent frame(s)" in status


def test_scene_chain_expands_compact_i2v_guide_mask_to_spatial_chunk_mask():
    sample_calls.clear()
    node = FunPackLTXAVSceneChainSampler()
    samples = torch.zeros(1, 2, 5, 24, 3)
    samples[:, :, 0] = 7.0
    mask = torch.ones(1, 1, 5, 1, 1)
    mask[:, :, 0] = 0.0
    latent_template = {"samples": samples, "noise_mask": mask}
    positive = [scene_cond(0), scene_cond(1)]

    _, status, scene_count, _ = node.sample(
        model=object(),
        vae=FakeVAE(),
        positive=positive,
        negative=[],
        sampler=object(),
        sigmas=torch.tensor([1.0, 0.0]),
        seed=45,
        latent_template=latent_template,
        num_frames_per_scene=5,
        frame_overlap=2,
        cfg=1.0,
        max_scenes=2,
        carry_i2v_guides=True,
    )

    second_call = sample_calls[1]
    assert scene_count == 2
    assert second_call["noise_mask"].shape == second_call["latent_image"].shape
    assert torch.all(second_call["noise_mask"][:, :, 5] == 0.0)
    assert "i2v guide tokens=1 latent frame(s)" in status


def test_scene_chain_does_not_carry_i2v_guides_by_default():
    inputs = FunPackLTXAVSceneChainSampler.INPUT_TYPES()["required"]["carry_i2v_guides"][1]
    assert inputs["default"] is False

    sample_calls.clear()
    node = FunPackLTXAVSceneChainSampler()
    samples = torch.zeros(1, 2, 5, 1, 1)
    samples[:, :, 0] = 7.0
    mask = torch.ones(1, 1, 5, 1, 1)
    mask[:, :, 0] = 0.0
    latent_template = {"samples": samples, "noise_mask": mask}
    positive = [scene_cond(0), scene_cond(1)]

    _, status, scene_count, _ = node.sample(
        model=object(),
        vae=FakeVAE(),
        positive=positive,
        negative=[],
        sampler=object(),
        sigmas=torch.tensor([1.0, 0.0]),
        seed=50,
        latent_template=latent_template,
        num_frames_per_scene=5,
        frame_overlap=2,
        cfg=1.0,
        max_scenes=2,
    )

    second_call = sample_calls[1]
    assert scene_count == 2
    assert torch.all(second_call["latent_image"][:, :, 2] == 0.0)
    assert torch.all(second_call["noise_mask"][:, :, :2] == 0.0)
    assert torch.all(second_call["noise_mask"][:, :, 2:] == 1.0)
    assert "i2v guide" not in status