# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

import json
from graphlib import TopologicalSorter
from pathlib import Path

import pytest
from comfyui_vllm_omni import nodes as omni_nodes

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]

WORKFLOW = (
    Path(__file__).resolve().parents[4]
    / "apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni MiniMax-H3 Reference to Video.json"
)


@pytest.fixture
def workflow():
    return json.loads(WORKFLOW.read_text())


def test_reference_workflow_connections(workflow):
    nodes = {node["id"]: node for node in workflow["nodes"]}
    links = {link[0]: link for link in workflow["links"]}
    assert len(nodes) == len(workflow["nodes"])
    assert len(links) == len(workflow["links"])
    graph: dict[int, set[int]] = {node_id: set() for node_id in nodes}
    for link_id, source, output_slot, target, input_slot, kind in links.values():
        output = nodes[source]["outputs"][output_slot]
        input_ = nodes[target]["inputs"][input_slot]
        assert output["type"] == input_["type"] == kind
        assert link_id in output["links"]
        assert input_["link"] == link_id
        graph[target].add(source)
    assert len(tuple(TopologicalSorter(graph).static_order())) == len(nodes)
    for node in nodes.values():
        for slot, input_ in enumerate(node.get("inputs", [])):
            if input_["link"] is not None:
                assert links[input_["link"]][3:5] == [node["id"], slot]
        for slot, output in enumerate(node.get("outputs", [])):
            for link_id in output.get("links") or []:
                assert links[link_id][1:3] == [node["id"], slot]


def test_reference_workflow_matches_omni_node_interfaces(workflow):
    for node in workflow["nodes"]:
        if not node["type"].startswith("VLLMOmni"):
            continue
        cls = getattr(omni_nodes, node["type"])
        schema = cls.INPUT_TYPES()
        inputs = {**schema.get("required", {}), **schema.get("optional", {})}
        for input_ in node["inputs"]:
            assert input_["type"] == inputs[input_["name"]][0]
        assert tuple(output["type"] for output in node["outputs"]) == cls.RETURN_TYPES
        expected_widgets = len(schema.get("required", {}))
        if "seed" in schema.get("required", {}):
            expected_widgets += 1
        assert len(node["widgets_values"]) == expected_widgets
    refs = next(node for node in workflow["nodes"] if node["type"] == "VLLMOmniVideoReferences")
    assert [input_["name"] for input_ in refs["inputs"]] == list(
        omni_nodes.VLLMOmniVideoReferences.INPUT_TYPES()["optional"]
    )


def test_reference_workflow_default_route_and_output(workflow):
    nodes = {node["id"]: node for node in workflow["nodes"]}
    links = {link[0]: link for link in workflow["links"]}
    generate = next(node for node in nodes.values() if node["type"] == "VLLMOmniGenerateVideo")
    inputs = {input_["name"]: input_["link"] for input_ in generate["inputs"]}
    assert inputs["frame"] is None
    assert inputs["lora"] is None
    for name, kind in (
        ("references", "VLLMOmniVideoReferences"),
        ("sampling_params", "VLLMOmniDiffusionSampling"),
        ("model_params", "VLLMOmniMiniMaxH3Params"),
    ):
        assert nodes[links[inputs[name]][1]]["type"] == kind
    save = next(node for node in nodes.values() if node["type"] == "SaveVideo")
    assert links[save["inputs"][0]["link"]][1] == generate["id"]
    assert save["widgets_values"] == ["video/MiniMax-H3-Ref2VA", "mp4", "h264"]


def test_reference_workflow_defaults_and_portable_assets(workflow):
    by_type = {node["type"]: node for node in workflow["nodes"]}
    values = by_type["VLLMOmniGenerateVideo"]["widgets_values"]
    assert values[0] == "http://localhost:8000/v1"
    assert values[1] == "MiniMaxAI/MiniMax-H3"
    assert "<Picture 1>" in values[2]
    assert values[4:8] == [1344, 768, 24, 5.167]
    num_frames = round(values[7] * values[6])
    assert num_frames == 124
    assert values[3] == ""
    assert (num_frames - 5) % 17 == 0
    assert 4 <= values[7] <= 15
    assert by_type["VLLMOmniDiffusionSampling"]["widgets_values"] == [1, 50, 1.0, 1.0, False, False, 42, "fixed"]
    assert by_type["VLLMOmniMiniMaxH3Params"]["widgets_values"] == [3.0, 12.0]
    assert by_type["VLLMOmniRemoteLoRA"]["widgets_values"][0] == ""
    for kind in ("LoadImage", "LoadVideo", "LoadAudio"):
        filename = by_type[kind]["widgets_values"][0]
        assert filename and Path(filename).name == filename
        assert ":" not in filename and "\\" not in filename
    assert set(by_type) == {
        "LoadImage",
        "LoadVideo",
        "LoadAudio",
        "SaveVideo",
        "MarkdownNote",
        "VLLMOmniVideoReferences",
        "VLLMOmniGenerateVideo",
        "VLLMOmniDiffusionSampling",
        "VLLMOmniMiniMaxH3Params",
        "VLLMOmniRemoteLoRA",
    }