from smart_metadata_reader.graph import GraphIndex


def select_sampler(prompt):
    from smart_metadata_reader.sampler_selector import select_final_sampler

    return select_final_sampler(GraphIndex(prompt=prompt, workflow=None))


def test_selects_ksampler_upstream_of_save_image_via_vae_decode():
    prompt = {
        "1": {"class_type": "KSampler", "inputs": {"positive": ["9", 0]}},
        "2": {"class_type": "VAEDecode", "inputs": {"samples": ["1", 0]}},
        "3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
        "9": {"class_type": "CLIPTextEncode", "inputs": {"text": "positive"}},
    }

    selection = select_sampler(prompt)

    assert selection.sampler.node_id == "1"
    assert selection.confidence == 1.0
    assert "SaveImage 3.images <- VAEDecode 2.samples <- KSampler 1" in selection.debug_trace


def test_selects_ksampler_upstream_of_preview_image_via_vae_decode():
    prompt = {
        "1": {"class_type": "KSampler", "inputs": {"positive": ["9", 0]}},
        "2": {"class_type": "VAEDecode", "inputs": {"samples": ["1", 0]}},
        "3": {"class_type": "PreviewImage", "inputs": {"images": ["2", 0]}},
        "9": {"class_type": "CLIPTextEncode", "inputs": {"text": "positive"}},
    }

    selection = select_sampler(prompt)

    assert selection.sampler.node_id == "1"
    assert "PreviewImage 3.images" in selection.debug_trace


def test_selects_sampler_through_image_passthrough_before_save_image():
    prompt = {
        "1": {"class_type": "KSampler", "inputs": {"positive": ["9", 0]}},
        "2": {"class_type": "VAEDecode", "inputs": {"samples": ["1", 0]}},
        "3": {"class_type": "ImagePassThrough", "inputs": {"image": ["2", 0]}},
        "4": {"class_type": "SaveImage", "inputs": {"images": ["3", 0]}},
        "9": {"class_type": "CLIPTextEncode", "inputs": {"text": "positive"}},
    }

    selection = select_sampler(prompt)

    assert selection.sampler.node_id == "1"
    assert "ImagePassThrough 3.image" in selection.debug_trace


def test_multiple_samplers_selects_one_connected_to_final_save_image():
    prompt = {
        "1": {"class_type": "KSampler", "inputs": {"positive": ["9", 0]}},
        "2": {"class_type": "VAEDecode", "inputs": {"samples": ["1", 0]}},
        "3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
        "8": {"class_type": "KSampler", "inputs": {"positive": ["9", 0]}},
        "9": {"class_type": "CLIPTextEncode", "inputs": {"text": "positive"}},
    }

    selection = select_sampler(prompt)

    assert selection.sampler.node_id == "1"
    assert "Selected KSampler 1 because it is upstream of final SaveImage" in selection.debug_trace


def test_fallback_selects_sampler_with_conditioning_inputs_when_no_output_node_exists():
    prompt = {
        "5": {"class_type": "PrimitiveString", "inputs": {"text": "not a sampler"}},
        "6": {
            "class_type": "KSampler",
            "inputs": {"positive": ["9", 0], "negative": ["10", 0]},
        },
        "9": {"class_type": "CLIPTextEncode", "inputs": {"text": "positive"}},
        "10": {"class_type": "CLIPTextEncode", "inputs": {"text": "negative"}},
    }

    selection = select_sampler(prompt)

    assert selection.sampler.node_id == "6"
    assert selection.confidence < 1.0
    assert "sampler selected by fallback heuristic" in selection.debug_trace


def test_extracts_ksampler_settings():
    from smart_metadata_reader.adapters.samplers import extract_sampler_settings
    from smart_metadata_reader.models import NodeRecord

    sampler = NodeRecord(
        node_id="1",
        class_type="KSampler",
        inputs={
            "seed": 123,
            "steps": 30,
            "cfg": 7.5,
            "sampler_name": "euler",
            "scheduler": "normal",
            "denoise": 0.85,
        },
    )

    assert extract_sampler_settings(sampler) == {
        "seed": 123,
        "steps": 30,
        "cfg": 7.5,
        "sampler_name": "euler",
        "scheduler": "normal",
        "denoise": 0.85,
    }


def test_extracts_ksampler_advanced_and_custom_field_variants():
    from smart_metadata_reader.adapters.samplers import extract_sampler_settings
    from smart_metadata_reader.models import NodeRecord

    advanced = NodeRecord(
        node_id="2",
        class_type="KSamplerAdvanced",
        inputs={
            "noise_seed": 456,
            "steps": 22,
            "cfg_scale": 6.0,
            "sampler": "dpmpp_2m",
            "schedule": "karras",
        },
    )
    custom = NodeRecord(
        node_id="3",
        class_type="SamplerCustom",
        inputs={
            "noise_seed": 789,
            "steps": 18,
            "cfg": 5.5,
            "sampler_name": "custom_sampler",
            "scheduler": "sgm_uniform",
        },
    )

    assert extract_sampler_settings(advanced) == {
        "seed": 456,
        "steps": 22,
        "cfg": 6.0,
        "sampler_name": "dpmpp_2m",
        "scheduler": "karras",
    }
    assert extract_sampler_settings(custom) == {
        "seed": 789,
        "steps": 18,
        "cfg": 5.5,
        "sampler_name": "custom_sampler",
        "scheduler": "sgm_uniform",
    }