import json
from pathlib import Path


WORKFLOW = (
    Path(__file__).parents[1]
    / "example_workflows"
    / "MiniMax H3 Per-Row Video Inpaint.json"
)


def _load_workflow():
    return json.loads(WORKFLOW.read_text(encoding="utf-8"))


def test_inpaint_workflow_links_are_consistent():
    workflow = _load_workflow()
    nodes = {node["id"]: node for node in workflow["nodes"]}
    link_ids = set()

    for link in workflow["links"]:
        link_id, origin_id, origin_slot, target_id, target_slot, link_type = link
        assert link_id not in link_ids
        link_ids.add(link_id)
        assert origin_id in nodes
        assert target_id in nodes
        output = nodes[origin_id]["outputs"][origin_slot]
        input_ = nodes[target_id]["inputs"][target_slot]
        assert link_id in output["links"]
        assert input_["link"] == link_id
        assert output["type"] == link_type
        assert input_["type"] in {link_type, "*"}

    assert workflow["last_link_id"] == max(link_ids)
    assert workflow["last_node_id"] == max(nodes)


def test_inpaint_workflow_contains_complete_h3_mask_path():
    workflow = _load_workflow()
    types = {node["type"] for node in workflow["nodes"]}
    assert {
        "MiniMaxH3TrimSourceAV",
        "MiniMaxH3MaskGridPreview",
        "MiniMaxH3SetGenerationMask",
        "MiniMaxH3SigmaShift",
        "MiniMaxH3PerRowMaskPatch",
        "MiniMaxH3ImageToVideo",
        "LTXVConcatAVLatent",
        "SamplerCustomAdvanced",
        "LTXVSeparateAVLatent",
    } <= types


def test_trimmed_length_controls_conditioning_and_source_latent_controls_sampler():
    workflow = _load_workflow()
    nodes = {node["id"]: node for node in workflow["nodes"]}
    links = {link[0]: link for link in workflow["links"]}

    trim = next(node for node in nodes.values() if node["type"] == "MiniMaxH3TrimSourceAV")
    conditioning = next(
        node for node in nodes.values() if node["type"] == "MiniMaxH3ImageToVideo"
    )
    sampler = next(node for node in nodes.values() if node["type"] == "SamplerCustomAdvanced")
    set_mask = next(
        node for node in nodes.values() if node["type"] == "MiniMaxH3SetGenerationMask"
    )

    length_link = links[conditioning["inputs"][7]["link"]]
    assert length_link[1:3] == [trim["id"], 2]

    latent_link = links[sampler["inputs"][4]["link"]]
    assert latent_link[1:3] == [set_mask["id"], 0]
    assert conditioning["outputs"][1]["links"] is None