import re

import dagster as dg
import pytest

from dagster_tests.core_tests.selector_tests.test_subset_selector import (
    foo_job,
    get_asset_selection_job,
)


def test_subset_for_execution():
    recon_job = dg.reconstructable(foo_job)
    sub_job = recon_job.get_subset(op_selection=["*add_nums"])

    assert sub_job.op_selection == {"*add_nums"}

    with dg.instance_for_test() as instance:
        result = dg.execute_job(sub_job, instance)
        assert result.success
        assert set([event.step_key for event in result.all_events if event.is_step_event]) == {
            "add_nums",
            "return_one",
            "return_two",
        }


def test_asset_subset_for_execution():
    recon_job = dg.reconstructable(get_asset_selection_job)
    sub_job = recon_job.get_subset(
        op_selection=None, asset_selection=frozenset({dg.AssetKey("my_asset")})
    )
    assert sub_job.asset_selection == {dg.AssetKey("my_asset")}

    with dg.instance_for_test() as instance:
        result = dg.execute_job(sub_job, instance, asset_selection=[dg.AssetKey("my_asset")])
        assert result.success

        materializations = result.filter_events(lambda evt: evt.is_step_materialization)
        assert len(materializations) == 1
        assert materializations[0].asset_key == dg.AssetKey("my_asset")


def test_reexecute_asset_subset():
    with dg.instance_for_test() as instance:
        result = dg.execute_job(
            dg.reconstructable(get_asset_selection_job),
            instance,
            asset_selection=[dg.AssetKey("my_asset")],
        )
        assert result.success
        materializations = [event for event in result.all_events if event.is_step_materialization]
        assert len(materializations) == 1
        assert materializations[0].asset_key == dg.AssetKey("my_asset")

        run = instance.get_run_by_id(result.run_id)
        assert run.asset_selection == {dg.AssetKey("my_asset")}  # pyright: ignore[reportOptionalMemberAccess]

        reexecution_result = dg.execute_job(
            dg.reconstructable(get_asset_selection_job),
            instance,
            reexecution_options=dg.ReexecutionOptions(parent_run_id=result.run_id),
        )

        assert reexecution_result.success
        materializations = reexecution_result.filter_events(lambda evt: evt.is_step_materialization)
        assert len(materializations) == 1
        assert materializations[0].asset_key == dg.AssetKey("my_asset")
        run = instance.get_run_by_id(reexecution_result.run_id)
        assert run.asset_selection == {dg.AssetKey("my_asset")}  # pyright: ignore[reportOptionalMemberAccess]


def test_execute_job_with_op_selection_single_clause():
    with dg.instance_for_test() as instance:
        with dg.execute_job(
            dg.reconstructable(foo_job),
            instance,
        ) as result_full:
            assert result_full.success
            assert result_full.output_for_node("add_one") == 7
            assert len(result_full.get_step_success_events()) == 5

        with dg.execute_job(
            dg.reconstructable(foo_job), op_selection=["*add_nums"], instance=instance
        ) as result_up:
            assert result_up.success
            assert result_up.output_for_node("add_nums") == 3
            assert len(result_up.get_step_success_events()) == 3

        with dg.execute_job(
            dg.reconstructable(foo_job),
            instance,
            run_config={
                "ops": {"add_nums": {"inputs": {"num1": {"value": 1}, "num2": {"value": 2}}}}
            },
            op_selection=["add_nums++"],
        ) as job_result_down:
            assert job_result_down.success
            assert job_result_down.output_for_node("add_one") == 7
            assert len(job_result_down.get_step_success_events()) == 3


def test_execute_job_with_op_selection_multi_clauses():
    with dg.instance_for_test() as instance:
        with dg.execute_job(
            dg.reconstructable(foo_job),
            instance,
            op_selection=["return_one", "return_two", "add_nums+"],
        ) as result_multi_disjoint:
            assert result_multi_disjoint.success
            assert result_multi_disjoint.output_for_node("multiply_two") == 6
            assert len(result_multi_disjoint.get_step_success_events()) == 4

        with dg.execute_job(
            dg.reconstructable(foo_job),
            instance,
            op_selection=["return_one++", "add_nums+", "return_two"],
        ) as result_multi_overlap:
            assert result_multi_overlap.success
            assert result_multi_overlap.output_for_node("multiply_two") == 6
            assert len(result_multi_overlap.get_step_success_events()) == 4

        with pytest.raises(
            dg.DagsterInvalidSubsetError,
            match=re.escape("No qualified ops to execute found for op_selection"),
        ):
            dg.execute_job(dg.reconstructable(foo_job), instance, op_selection=["a", "*add_nums"])


def test_execute_job_with_op_selection_invalid():
    invalid_input = ["return_one,return_two"]

    with dg.instance_for_test() as instance:
        with pytest.raises(
            dg.DagsterInvalidSubsetError,
            match=re.escape(f"No qualified ops to execute found for op_selection={invalid_input}"),
        ):
            dg.execute_job(
                dg.reconstructable(foo_job), op_selection=invalid_input, instance=instance
            )


def test_reexecute_job_with_step_selection_single_clause():
    with dg.instance_for_test() as instance:
        with dg.execute_job(dg.reconstructable(foo_job), instance=instance) as result_full:
            assert result_full.success
            assert result_full.output_for_node("add_one") == 7
            assert len(result_full.get_step_success_events()) == 5

            with dg.execute_job(
                dg.reconstructable(foo_job),
                instance=instance,
                reexecution_options=dg.ReexecutionOptions(parent_run_id=result_full.run_id),
            ) as reexecution_result_full:
                assert reexecution_result_full.success
                assert len(reexecution_result_full.get_step_success_events()) == 5
                assert reexecution_result_full.output_for_node("add_one") == 7

            with dg.execute_job(
                dg.reconstructable(foo_job),
                instance=instance,
                reexecution_options=dg.ReexecutionOptions(
                    parent_run_id=result_full.run_id,
                    step_selection=["*add_nums"],
                ),
            ) as reexecution_result_up:
                assert reexecution_result_up.success
                assert reexecution_result_up.output_for_node("add_nums") == 3

            with dg.execute_job(
                dg.reconstructable(foo_job),
                instance=instance,
                reexecution_options=dg.ReexecutionOptions(
                    parent_run_id=result_full.run_id,
                    step_selection=["add_nums++"],
                ),
                raise_on_error=True,
            ) as reexecution_result_down:
                assert reexecution_result_down.success
                assert reexecution_result_down.output_for_node("add_one") == 7


def test_reexecute_job_with_step_selection_multi_clauses():
    with dg.instance_for_test() as instance:
        with dg.execute_job(dg.reconstructable(foo_job), instance=instance) as result_full:
            assert result_full.success
            assert result_full.output_for_node("add_one") == 7
            assert len(result_full.get_step_success_events()) == 5

        with dg.execute_job(
            dg.reconstructable(foo_job),
            instance=instance,
            reexecution_options=dg.ReexecutionOptions(
                parent_run_id=result_full.run_id,
                step_selection=["return_one", "return_two", "add_nums+"],
            ),
        ) as result_multi_disjoint:
            assert result_multi_disjoint.success
            assert result_multi_disjoint.output_for_node("multiply_two") == 6

        with dg.execute_job(
            dg.reconstructable(foo_job),
            instance=instance,
            reexecution_options=dg.ReexecutionOptions(
                parent_run_id=result_full.run_id,
                step_selection=["return_one++", "return_two", "add_nums+"],
            ),
        ) as result_multi_overlap:
            assert result_multi_overlap.success
            assert result_multi_overlap.output_for_node("multiply_two") == 6

        with pytest.raises(
            dg.DagsterExecutionStepNotFoundError,
            match="Step selection refers to unknown step: a",
        ):
            dg.execute_job(
                dg.reconstructable(foo_job),
                instance=instance,
                reexecution_options=dg.ReexecutionOptions(
                    parent_run_id=result_full.run_id,
                    step_selection=["a", "*add_nums"],
                ),
            )

        with pytest.raises(
            dg.DagsterExecutionStepNotFoundError,
            match="Step selection refers to unknown steps: a, b",
        ):
            dg.execute_job(
                dg.reconstructable(foo_job),
                instance=instance,
                reexecution_options=dg.ReexecutionOptions(
                    parent_run_id=result_full.run_id,
                    step_selection=["a+", "*b"],
                ),
            )
