import contextlib
import importlib.util
import sys
import types
from pathlib import Path

import pytest
import torch


def _load_node_package(monkeypatch, tmp_path):
    comfy = types.ModuleType("comfy")
    comfy.__path__ = []
    management = types.ModuleType("comfy.model_management")
    management.load_calls = []
    management.interrupt_checks = 0
    management.progress_totals = []
    management.progress_updates = []
    management.load_models_gpu = lambda models, force_full_load=False: (
        management.load_calls.append((models, force_full_load))
    )
    management.throw_exception_if_processing_interrupted = lambda: setattr(
        management, "interrupt_checks", management.interrupt_checks + 1
    )
    patcher = types.ModuleType("comfy.model_patcher")
    utils = types.ModuleType("comfy.utils")
    def progress_bar(total):
        management.progress_totals.append(total)
        return types.SimpleNamespace(
            update=lambda value: management.progress_updates.append(value)
        )

    utils.ProgressBar = progress_bar
    folder_paths = types.ModuleType("folder_paths")
    folder_paths.models_dir = str(tmp_path)
    folder_paths._paths = []
    folder_paths.add_model_folder_path = (
        lambda name, path, is_default=False: folder_paths._paths.append(path)
        if path not in folder_paths._paths
        else None
    )
    folder_paths.get_folder_paths = lambda name: folder_paths._paths

    monkeypatch.setitem(sys.modules, "comfy", comfy)
    monkeypatch.setitem(sys.modules, "comfy.model_management", management)
    monkeypatch.setitem(sys.modules, "comfy.model_patcher", patcher)
    monkeypatch.setitem(sys.modules, "comfy.utils", utils)
    monkeypatch.setitem(sys.modules, "folder_paths", folder_paths)

    root = Path(__file__).parents[1]
    package_name = f"comfyui_sam_audio_node_test_{tmp_path.name}"
    spec = importlib.util.spec_from_file_location(
        package_name,
        root / "__init__.py",
        submodule_search_locations=[str(root)],
    )
    module = importlib.util.module_from_spec(spec)
    monkeypatch.setitem(sys.modules, package_name, module)
    spec.loader.exec_module(module)
    return module, management


class FakeBatch:
    def __init__(self, audios, descriptions, anchors, masked_videos):
        self.audios = audios
        self.descriptions = descriptions
        self.anchors = anchors
        self.masked_videos = masked_videos

    def to(self, device):
        return self


class FakeProcessor:
    audio_sampling_rate = 8
    audio_hop_length = 2

    def __init__(self):
        self.calls = []

    def __call__(self, audios, descriptions, anchors=None, masked_videos=None):
        self.calls.append(
            {
                "audios": audios,
                "descriptions": descriptions,
                "anchors": anchors,
                "masked_videos": masked_videos,
            }
        )
        return FakeBatch(audios, descriptions, anchors, masked_videos)

    def mask_videos(self, videos, masks):
        return [
            video * mask.eq(0) for video, mask in zip(videos, masks, strict=False)
        ]


class FakeModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.calls = []

    def separate(self, batch, **kwargs):
        self.calls.append(kwargs)
        target = [row.mean(0) for row in batch.audios]
        return types.SimpleNamespace(target=target, residual=[-row for row in target])


def test_text_node_runs_full_audio_batch_with_standard_shapes(monkeypatch, tmp_path):
    package, management = _load_node_package(monkeypatch, tmp_path)
    nodes = sys.modules[f"{package.__name__}.nodes"]
    runtime = sys.modules[f"{package.__name__}.sam_audio_comfy.runtime"]
    pipeline = runtime.SAMAudioPipeline(
        model=FakeModel(),
        processor=FakeProcessor(),
        patcher=object(),
        device=torch.device("cpu"),
        source="fake",
    )

    target, residual = nodes.SAMAudioTextSeparator().separate(
        pipeline=pipeline,
        audio={"waveform": torch.ones(2, 2, 8), "sample_rate": 8},
        description="test sound",
        predict_spans=False,
        seed=7,
        inference_steps=32,
    )

    assert target["waveform"].shape == (2, 1, 8)
    assert residual["waveform"].shape == (2, 1, 8)
    assert target["sample_rate"] == 8
    loader_inputs = package.NODE_CLASS_MAPPINGS[
        "SAMAudioPipelineLoader"
    ].INPUT_TYPES()["required"]
    assert loader_inputs["attention_backend"][0] == ["pytorch", "comfy_kitchen"]
    sampler_inputs = package.NODE_CLASS_MAPPINGS[
        "SAMAudioTextSeparator"
    ].INPUT_TYPES()["required"]
    sampler_optional = package.NODE_CLASS_MAPPINGS[
        "SAMAudioTextSeparator"
    ].INPUT_TYPES()["optional"]
    loader_optional = package.NODE_CLASS_MAPPINGS[
        "SAMAudioPipelineLoader"
    ].INPUT_TYPES()["optional"]
    assert sampler_inputs["seed"][1]["control_after_generate"] is True
    assert sampler_inputs["chunk_duration"][1]["default"] == 10.0
    assert sampler_inputs["chunk_overlap"][1]["default"] == 1.0
    assert sampler_optional["reranking_candidates"][1]["default"] == 1
    assert loader_optional["precision"][0] == ["auto", "bf16", "fp32"]
    assert loader_optional["initial_device"][0] == ["cpu", "gpu"]
    assert set(package.NODE_CLASS_MAPPINGS) == {
        "SAMAudioPipelineLoader",
        "SAMAudioSpanPrompt",
        "SAMAudioTextSeparator",
        "SAMAudioSpanSeparator",
        "SAMAudioVisualSeparator",
        "SAMAudioVideoSeparator",
    }
    assert management.load_calls == [([pipeline.patcher], True)]
    assert management.progress_totals == [32]
    assert sum(management.progress_updates) == 32
    assert management.interrupt_checks == 34


def _fake_pipeline(runtime, processor=None):
    return runtime.SAMAudioPipeline(
        model=FakeModel(),
        processor=processor or FakeProcessor(),
        patcher=object(),
        device=torch.device("cpu"),
        source="fake",
    )


def test_text_node_chunks_crossfades_and_advances_seed(monkeypatch, tmp_path):
    package, _ = _load_node_package(monkeypatch, tmp_path)
    nodes = sys.modules[f"{package.__name__}.nodes"]
    runtime = sys.modules[f"{package.__name__}.sam_audio_comfy.runtime"]
    processor = FakeProcessor()
    pipeline = _fake_pipeline(runtime, processor)
    seeds = []

    @contextlib.contextmanager
    def record_seed(seed, device):
        seeds.append(seed)
        yield

    monkeypatch.setattr(nodes, "seeded_inference", record_seed)
    target, residual = nodes.SAMAudioTextSeparator().separate(
        pipeline=pipeline,
        audio={"waveform": torch.ones(2, 2, 18), "sample_rate": 8},
        description="test sound",
        predict_spans=False,
        seed=7,
        inference_steps=32,
        chunk_duration=0.75,
        chunk_overlap=0.25,
    )

    assert [call["audios"][0].shape[-1] for call in processor.calls] == [6, 6, 6, 6]
    assert seeds == [7, 8, 9, 10]
    assert torch.equal(target["waveform"], torch.ones(2, 1, 18))
    assert torch.equal(residual["waveform"], -torch.ones(2, 1, 18))


def test_span_node_localizes_global_spans_for_each_chunk(monkeypatch, tmp_path):
    package, _ = _load_node_package(monkeypatch, tmp_path)
    nodes = sys.modules[f"{package.__name__}.nodes"]
    runtime = sys.modules[f"{package.__name__}.sam_audio_comfy.runtime"]
    processor = FakeProcessor()
    pipeline = _fake_pipeline(runtime, processor)
    spans = (
        runtime.SpanPrompt("+", 0.5, 1.25),
        runtime.SpanPrompt("-", 1.5, 1.75),
    )

    nodes.SAMAudioSpanSeparator().separate(
        pipeline=pipeline,
        audio={"waveform": torch.ones(1, 1, 16), "sample_rate": 8},
        description="test sound",
        spans=spans,
        seed=0,
        inference_steps=32,
        chunk_duration=1.0,
        chunk_overlap=0.25,
    )

    anchors = [call["anchors"][0] for call in processor.calls]
    assert anchors[0] == [runtime.SpanPrompt("+", 0.5, 1.0)]
    assert anchors[1] == [
        runtime.SpanPrompt("+", 0.0, 0.5),
        runtime.SpanPrompt("-", 0.75, 1.0),
    ]
    assert anchors[2] == [runtime.SpanPrompt("-", 0.0, 0.25)]


def test_visual_node_slices_frames_to_each_audio_chunk(monkeypatch, tmp_path):
    package, _ = _load_node_package(monkeypatch, tmp_path)
    nodes = sys.modules[f"{package.__name__}.nodes"]
    runtime = sys.modules[f"{package.__name__}.sam_audio_comfy.runtime"]
    processor = FakeProcessor()
    pipeline = _fake_pipeline(runtime, processor)
    images = torch.linspace(0, 1, 8).view(8, 1, 1, 1).expand(-1, 2, 2, 3)

    nodes.SAMAudioVisualSeparator().separate(
        pipeline=pipeline,
        audio={"waveform": torch.ones(1, 1, 16), "sample_rate": 8},
        images=images,
        mask=torch.zeros(2, 2),
        description="",
        seed=0,
        inference_steps=32,
        chunk_duration=1.0,
        chunk_overlap=0.25,
    )

    frame_counts = [call["masked_videos"][0].shape[0] for call in processor.calls]
    assert frame_counts == [4, 4, 2]


def test_video_node_uses_native_comfy_video_frames_and_audio(monkeypatch, tmp_path):
    package, _ = _load_node_package(monkeypatch, tmp_path)
    nodes = sys.modules[f"{package.__name__}.nodes"]
    runtime = sys.modules[f"{package.__name__}.sam_audio_comfy.runtime"]
    processor = FakeProcessor()
    pipeline = _fake_pipeline(runtime, processor)
    images = torch.ones(8, 2, 2, 3)
    audio = {"waveform": torch.ones(1, 1, 16), "sample_rate": 8}
    video = types.SimpleNamespace(
        get_components=lambda: types.SimpleNamespace(
            images=images,
            audio=audio,
            frame_rate=4,
        )
    )

    target, residual = nodes.SAMAudioVideoSeparator().separate(
        pipeline=pipeline,
        video=video,
        mask=torch.ones(2, 2),
        description="",
        seed=0,
        inference_steps=32,
        chunk_duration=0,
        chunk_overlap=0,
    )

    assert processor.calls[0]["audios"][0].shape == (1, 16)
    assert processor.calls[0]["masked_videos"][0].shape == (8, 3, 2, 2)
    assert target["waveform"].shape == (1, 1, 16)
    assert residual["waveform"].shape == (1, 1, 16)
    video_inputs = package.NODE_CLASS_MAPPINGS[
        "SAMAudioVideoSeparator"
    ].INPUT_TYPES()["required"]
    assert video_inputs["video"][0] == "VIDEO"


def test_video_node_trims_audio_to_native_frame_duration(monkeypatch, tmp_path):
    package, _ = _load_node_package(monkeypatch, tmp_path)
    nodes = sys.modules[f"{package.__name__}.nodes"]
    runtime = sys.modules[f"{package.__name__}.sam_audio_comfy.runtime"]
    processor = FakeProcessor()
    pipeline = _fake_pipeline(runtime, processor)
    video = types.SimpleNamespace(
        get_components=lambda: types.SimpleNamespace(
            images=torch.ones(8, 2, 2, 3),
            audio={"waveform": torch.ones(1, 1, 16), "sample_rate": 8},
            frame_rate=8,
        )
    )

    target, _ = nodes.SAMAudioVideoSeparator().separate(
        pipeline=pipeline,
        video=video,
        mask=torch.zeros(2, 2),
        description="",
        seed=0,
        inference_steps=2,
        chunk_duration=0,
        chunk_overlap=0,
    )

    assert processor.calls[0]["audios"][0].shape[-1] == 8
    assert target["waveform"].shape[-1] == 8


def test_reranking_is_lazy_and_passed_to_upstream(monkeypatch, tmp_path):
    package, _ = _load_node_package(monkeypatch, tmp_path)
    nodes = sys.modules[f"{package.__name__}.nodes"]
    runtime = sys.modules[f"{package.__name__}.sam_audio_comfy.runtime"]
    pipeline = _fake_pipeline(runtime)
    loaded = []
    pipeline.ensure_judge_ranker = lambda management: loaded.append(management)

    nodes.SAMAudioTextSeparator().separate(
        pipeline=pipeline,
        audio={"waveform": torch.ones(1, 1, 8), "sample_rate": 8},
        description="test sound",
        predict_spans=False,
        seed=0,
        inference_steps=2,
        reranking_candidates=4,
    )

    assert len(loaded) == 1
    assert pipeline.model.calls[0]["reranking_candidates"] == 4


def test_transformer_steps_report_progress_and_check_cancellation(monkeypatch, tmp_path):
    package, management = _load_node_package(monkeypatch, tmp_path)
    nodes = sys.modules[f"{package.__name__}.nodes"]
    runtime = sys.modules[f"{package.__name__}.sam_audio_comfy.runtime"]

    class SteppingModel(FakeModel):
        def __init__(self):
            super().__init__()
            self.transformer = torch.nn.Identity()

        def separate(self, batch, **kwargs):
            self.transformer(torch.zeros(1))
            self.transformer(torch.zeros(1))
            return super().separate(batch, **kwargs)

    pipeline = runtime.SAMAudioPipeline(
        model=SteppingModel(),
        processor=FakeProcessor(),
        patcher=object(),
        device=torch.device("cpu"),
        source="fake",
    )

    nodes.SAMAudioTextSeparator().separate(
        pipeline=pipeline,
        audio={"waveform": torch.ones(1, 1, 8), "sample_rate": 8},
        description="test sound",
        predict_spans=False,
        seed=0,
        inference_steps=2,
    )

    assert management.progress_totals == [2]
    assert sum(management.progress_updates) == 2
    assert management.interrupt_checks == 4
    assert not pipeline.model.transformer._forward_pre_hooks


def test_video_node_requires_embedded_audio(monkeypatch, tmp_path):
    package, _ = _load_node_package(monkeypatch, tmp_path)
    nodes = sys.modules[f"{package.__name__}.nodes"]
    runtime = sys.modules[f"{package.__name__}.sam_audio_comfy.runtime"]
    pipeline = _fake_pipeline(runtime)
    video = types.SimpleNamespace(
        get_components=lambda: types.SimpleNamespace(
            images=torch.ones(1, 2, 2, 3),
            audio=None,
            frame_rate=1,
        )
    )

    with pytest.raises(ValueError, match="has no audio track"):
        nodes.SAMAudioVideoSeparator().separate(
            pipeline=pipeline,
            video=video,
            mask=torch.ones(2, 2),
            description="",
            seed=0,
            inference_steps=32,
        )