import pytest

from dagster import (
    AssetMaterialization,
    CompositeSolidDefinition,
    DagsterInstance,
    DagsterType,
    DagsterUnknownResourceError,
    InputDefinition,
    ModeDefinition,
    OutputDefinition,
    ResourceDefinition,
    String,
    composite_solid,
    dagster_type_loader,
    dagster_type_materializer,
    execute_pipeline,
    pipeline,
    resource,
    solid,
    usable_as_dagster_type,
)
from dagster.core.definitions.executable import InMemoryExecutablePipeline
from dagster.core.execution.api import execute_run
from dagster.core.storage.type_storage import TypeStoragePlugin
from dagster.core.types.dagster_type import create_any_type


def get_resource_init_pipeline(resources_initted):
    @resource
    def resource_a(_):
        resources_initted['a'] = True
        yield 'A'

    @resource
    def resource_b(_):
        resources_initted['b'] = True
        yield 'B'

    @solid(required_resource_keys={'a'})
    def consumes_resource_a(context):
        assert context.resources.a == 'A'

    @solid(required_resource_keys={'b'})
    def consumes_resource_b(context):
        assert context.resources.b == 'B'

    @pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a, 'b': resource_b,})],)
    def selective_init_test_pipeline():
        consumes_resource_a()
        consumes_resource_b()

    return selective_init_test_pipeline


def test_filter_out_resources():
    @solid(required_resource_keys={'a'})
    def requires_resource_a(context):
        assert context.resources.a
        assert not hasattr(context.resources, 'b')

    @solid(required_resource_keys={'b'})
    def requires_resource_b(context):
        assert not hasattr(context.resources, 'a')
        assert context.resources.b

    @solid
    def not_resources(context):
        assert not hasattr(context.resources, 'a')
        assert not hasattr(context.resources, 'b')

    @pipeline(
        mode_defs=[
            ModeDefinition(
                resource_defs={
                    'a': ResourceDefinition.hardcoded_resource('foo'),
                    'b': ResourceDefinition.hardcoded_resource('bar'),
                }
            )
        ],
    )
    def room_of_requirement():
        requires_resource_a()
        requires_resource_b()
        not_resources()

    execute_pipeline(room_of_requirement)


def test_selective_init_resources():
    resources_initted = {}

    assert execute_pipeline(get_resource_init_pipeline(resources_initted)).success

    assert set(resources_initted.keys()) == {'a', 'b'}


def test_selective_init_resources_only_a():
    resources_initted = {}

    @resource
    def resource_a(_):
        resources_initted['a'] = True
        yield 'A'

    @resource
    def resource_b(_):
        resources_initted['b'] = True
        yield 'B'

    @solid(required_resource_keys={'a'})
    def consumes_resource_a(context):
        assert context.resources.a == 'A'

    @pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a, 'b': resource_b})])
    def selective_init_test_pipeline():
        consumes_resource_a()

    assert execute_pipeline(selective_init_test_pipeline).success

    assert set(resources_initted.keys()) == {'a'}


def test_execution_plan_subset_strict_resources():
    resources_initted = {}

    instance = DagsterInstance.ephemeral()

    pipeline_def = get_resource_init_pipeline(resources_initted)

    pipeline_run = instance.create_run_for_pipeline(
        pipeline_def, step_keys_to_execute=['consumes_resource_b.compute'],
    )

    result = execute_run(InMemoryExecutablePipeline(pipeline_def), pipeline_run, instance)

    assert result.success

    assert set(resources_initted.keys()) == {'b'}


def test_solid_selection_strict_resources():
    resources_initted = {}

    selective_init_test_pipeline = get_resource_init_pipeline(resources_initted)

    result = execute_pipeline(
        selective_init_test_pipeline.get_pipeline_subset_def({'consumes_resource_b'})
    )
    assert result.success

    assert set(resources_initted.keys()) == {'b'}


def test_solid_selection_with_aliases_strict_resources():
    resources_initted = {}

    @resource
    def resource_a(_):
        resources_initted['a'] = True
        yield 'A'

    @resource
    def resource_b(_):
        resources_initted['b'] = True
        yield 'B'

    @solid(required_resource_keys={'a'})
    def consumes_resource_a(context):
        assert context.resources.a == 'A'

    @solid(required_resource_keys={'b'})
    def consumes_resource_b(context):
        assert context.resources.b == 'B'

    @pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a, 'b': resource_b,})],)
    def selective_init_test_pipeline():
        consumes_resource_a.alias('alias_for_a')()
        consumes_resource_b()

    result = execute_pipeline(selective_init_test_pipeline.get_pipeline_subset_def({'alias_for_a'}))
    assert result.success

    assert set(resources_initted.keys()) == {'a'}


def create_composite_solid_pipeline(resources_initted):
    @resource
    def resource_a(_):
        resources_initted['a'] = True
        yield 'a'

    @resource
    def resource_b(_):
        resources_initted['b'] = True
        yield 'B'

    @solid(required_resource_keys={'a'})
    def consumes_resource_a(context):
        assert context.resources.a == 'A'

    @solid(required_resource_keys={'b'})
    def consumes_resource_b(context):
        assert context.resources.b == 'B'

    @solid
    def consumes_resource_b_error(context):
        assert context.resources.b == 'B'

    @composite_solid
    def wraps_a():
        consumes_resource_a()

    @composite_solid
    def wraps_b():
        consumes_resource_b()

    @composite_solid
    def wraps_b_error():
        consumes_resource_b()
        consumes_resource_b_error()

    @pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a, 'b': resource_b,})],)
    def selective_init_composite_test_pipeline():
        wraps_a()
        wraps_b()
        wraps_b_error()

    return selective_init_composite_test_pipeline


def test_solid_selection_strict_resources_within_composite():
    resources_initted = {}

    result = execute_pipeline(
        create_composite_solid_pipeline(resources_initted).get_pipeline_subset_def({'wraps_b'})
    )
    assert result.success

    assert set(resources_initted.keys()) == {'b'}


def test_execution_plan_subset_strict_resources_within_composite():
    resources_initted = {}

    pipeline_def = create_composite_solid_pipeline(resources_initted)

    instance = DagsterInstance.ephemeral()

    pipeline_run = instance.create_run_for_pipeline(
        pipeline_def, step_keys_to_execute=['wraps_b.consumes_resource_b.compute'],
    )

    result = execute_run(InMemoryExecutablePipeline(pipeline_def), pipeline_run, instance)

    assert result.success

    assert set(resources_initted.keys()) == {'b'}


def test_unknown_resource_composite_error():
    resources_initted = {}

    with pytest.raises(DagsterUnknownResourceError):
        execute_pipeline(
            create_composite_solid_pipeline(resources_initted).get_pipeline_subset_def(
                {'wraps_b_error'}
            )
        )


def test_execution_plan_subset_with_aliases():
    resources_initted = {}

    @resource
    def resource_a(_):
        resources_initted['a'] = True
        yield 'A'

    @resource
    def resource_b(_):
        resources_initted['b'] = True
        yield 'B'

    @solid(required_resource_keys={'a'})
    def consumes_resource_a(context):
        assert context.resources.a == 'A'

    @solid(required_resource_keys={'b'})
    def consumes_resource_b(context):
        assert context.resources.b == 'B'

    @pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a, 'b': resource_b,})],)
    def selective_init_test_pipeline_with_alias():
        consumes_resource_a()
        consumes_resource_b.alias('b_alias')()

    instance = DagsterInstance.ephemeral()

    pipeline_run = instance.create_run_for_pipeline(
        selective_init_test_pipeline_with_alias, step_keys_to_execute=['b_alias.compute'],
    )

    result = execute_run(
        InMemoryExecutablePipeline(selective_init_test_pipeline_with_alias), pipeline_run, instance
    )

    assert result.success

    assert set(resources_initted.keys()) == {'b'}


def test_custom_type_with_resource_dependent_hydration():
    def define_input_hydration_pipeline(should_require_resources):
        @resource
        def resource_a(_):
            yield 'A'

        @dagster_type_loader(
            String, required_resource_keys={'a'} if should_require_resources else set()
        )
        def InputHydration(context, hello):
            assert context.resources.a == 'A'
            return CustomType(hello)

        @usable_as_dagster_type(loader=InputHydration)
        class CustomType(str):
            pass

        @solid(input_defs=[InputDefinition('custom_type', CustomType)])
        def input_hydration_solid(context, custom_type):
            context.log.info(custom_type)

        @pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a})])
        def input_hydration_pipeline():
            input_hydration_solid()

        return input_hydration_pipeline

    under_required_pipeline = define_input_hydration_pipeline(should_require_resources=False)
    with pytest.raises(DagsterUnknownResourceError):
        execute_pipeline(
            under_required_pipeline,
            {'solids': {'input_hydration_solid': {'inputs': {'custom_type': 'hello'}}}},
        )

    sufficiently_required_pipeline = define_input_hydration_pipeline(should_require_resources=True)
    assert execute_pipeline(
        sufficiently_required_pipeline,
        {'solids': {'input_hydration_solid': {'inputs': {'custom_type': 'hello'}}}},
    ).success


def test_resource_dependent_hydration_with_selective_init():
    def get_resource_init_input_hydration_pipeline(resources_initted):
        @resource
        def resource_a(_):
            resources_initted['a'] = True
            yield 'A'

        @dagster_type_loader(String, required_resource_keys={'a'})
        def InputHydration(context, hello):
            assert context.resources.a == 'A'
            return CustomType(hello)

        @usable_as_dagster_type(loader=InputHydration)
        class CustomType(str):
            pass

        @solid(input_defs=[InputDefinition('custom_type', CustomType)])
        def input_hydration_solid(context, custom_type):
            context.log.info(custom_type)

        @solid(output_defs=[OutputDefinition(CustomType)])
        def source_custom_type(_):
            return CustomType('from solid')

        @pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a})])
        def selective_pipeline():
            input_hydration_solid(source_custom_type())

        return selective_pipeline

    resources_initted = {}
    assert execute_pipeline(get_resource_init_input_hydration_pipeline(resources_initted)).success
    assert set(resources_initted.keys()) == set()


def define_plugin_pipeline(
    should_require_resources=True, resources_initted=None, compatible_storage=True
):
    if resources_initted is None:
        resources_initted = {}

    class CustomStoragePlugin(TypeStoragePlugin):  # pylint: disable=no-init
        @classmethod
        def compatible_with_storage_def(cls, system_storage_def):
            return compatible_storage

        @classmethod
        def set_intermediate_object(
            cls, intermediate_storage, context, dagster_type, step_output_handle, value
        ):
            assert context.resources.a == 'A'
            return intermediate_storage.set_intermediate_object(
                dagster_type, step_output_handle, value
            )

        @classmethod
        def get_intermediate_object(
            cls, intermediate_storage, context, dagster_type, step_output_handle
        ):
            assert context.resources.a == 'A'
            return intermediate_storage.get_intermediate_object(dagster_type, step_output_handle)

        @classmethod
        def required_resource_keys(cls):
            return {'a'} if should_require_resources else set()

    @resource
    def resource_a(_):
        resources_initted['a'] = True
        yield 'A'

    CustomDagsterType = create_any_type(name='CustomType', auto_plugins=[CustomStoragePlugin])

    @solid(output_defs=[OutputDefinition(CustomDagsterType)])
    def output_solid(_context):
        return 'hello'

    @pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a})])
    def plugin_pipeline():
        output_solid()

    return plugin_pipeline


def test_custom_type_with_resource_dependent_storage_plugin():
    under_required_pipeline = define_plugin_pipeline(should_require_resources=False)
    with pytest.raises(DagsterUnknownResourceError):
        execute_pipeline(under_required_pipeline, {'storage': {'filesystem': {}}})

    resources_initted = {}
    sufficiently_required_pipeline = define_plugin_pipeline(
        should_require_resources=True, resources_initted=resources_initted
    )
    assert execute_pipeline(sufficiently_required_pipeline, {'storage': {'filesystem': {}}}).success
    assert set(resources_initted.keys()) == set('a')

    resources_initted = {}
    assert execute_pipeline(
        define_plugin_pipeline(resources_initted, compatible_storage=False)
    ).success
    assert set(resources_initted.keys()) == set()


def define_materialization_pipeline(should_require_resources=True, resources_initted=None):
    if resources_initted is None:
        resources_initted = {}

    @resource
    def resource_a(_):
        resources_initted['a'] = True
        yield 'A'

    @dagster_type_materializer(
        String, required_resource_keys={'a'} if should_require_resources else set()
    )
    def materialize(context, *_args, **_kwargs):
        assert context.resources.a == 'A'
        return AssetMaterialization('hello')

    CustomDagsterType = create_any_type(name='CustomType', materializer=materialize)

    @solid(output_defs=[OutputDefinition(CustomDagsterType)])
    def output_solid(_context):
        return 'hello'

    @pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a})])
    def output_pipeline():
        output_solid()

    return output_pipeline


def test_custom_type_with_resource_dependent_materialization():
    under_required_pipeline = define_materialization_pipeline(should_require_resources=False)
    with pytest.raises(DagsterUnknownResourceError):
        execute_pipeline(
            under_required_pipeline,
            {'solids': {'output_solid': {'outputs': [{'result': 'hello'}]}}},
        )

    resources_initted = {}
    sufficiently_required_pipeline = define_materialization_pipeline(
        should_require_resources=True, resources_initted=resources_initted
    )
    res = execute_pipeline(
        sufficiently_required_pipeline,
        {'solids': {'output_solid': {'outputs': [{'result': 'hello'}]}}},
    )
    assert res.success
    assert res.result_for_solid('output_solid').output_value() == 'hello'
    assert set(resources_initted.keys()) == set('a')

    resources_initted = {}
    assert execute_pipeline(
        define_materialization_pipeline(resources_initted=resources_initted)
    ).success
    assert set(resources_initted.keys()) == set()


def define_composite_materialization_pipeline(
    should_require_resources=True, resources_initted=None
):
    if resources_initted is None:
        resources_initted = {}

    @resource
    def resource_a(_):
        resources_initted['a'] = True
        yield 'A'

    @dagster_type_materializer(
        String, required_resource_keys={'a'} if should_require_resources else set()
    )
    def materialize(context, *_args, **_kwargs):
        assert context.resources.a == 'A'
        return AssetMaterialization('hello')

    CustomDagsterType = create_any_type(name='CustomType', materializer=materialize)

    @solid(output_defs=[OutputDefinition(CustomDagsterType)])
    def output_solid(_context):
        return 'hello'

    wrap_solid = CompositeSolidDefinition(
        name="wrap_solid",
        solid_defs=[output_solid],
        output_mappings=[OutputDefinition(CustomDagsterType).mapping_from('output_solid')],
    )

    @pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a})])
    def output_pipeline():
        wrap_solid()

    return output_pipeline


def test_custom_type_with_resource_dependent_composite_materialization():
    under_required_pipeline = define_composite_materialization_pipeline(
        should_require_resources=False
    )
    with pytest.raises(DagsterUnknownResourceError):
        execute_pipeline(
            under_required_pipeline, {'solids': {'wrap_solid': {'outputs': [{'result': 'hello'}]}}},
        )

    sufficiently_required_pipeline = define_composite_materialization_pipeline(
        should_require_resources=True
    )
    assert execute_pipeline(
        sufficiently_required_pipeline,
        {'solids': {'wrap_solid': {'outputs': [{'result': 'hello'}]}}},
    ).success

    # test that configured output materialization of the wrapping composite initializes resource
    resources_initted = {}
    assert execute_pipeline(
        define_composite_materialization_pipeline(resources_initted=resources_initted),
        {'solids': {'wrap_solid': {'outputs': [{'result': 'hello'}]}}},
    ).success
    assert set(resources_initted.keys()) == set('a')

    # test that configured output materialization of the inner solid initializes resource
    resources_initted = {}
    assert execute_pipeline(
        define_composite_materialization_pipeline(resources_initted=resources_initted),
        {
            'solids': {
                'wrap_solid': {'solids': {'output_solid': {'outputs': [{'result': 'hello'}]}}}
            }
        },
    ).success
    assert set(resources_initted.keys()) == set('a')

    # test that no output config will not initialize anything
    resources_initted = {}
    assert execute_pipeline(
        define_composite_materialization_pipeline(resources_initted=resources_initted),
    ).success
    assert set(resources_initted.keys()) == set()


def test_custom_type_with_resource_dependent_type_check():
    def define_type_check_pipeline(should_require_resources):
        @resource
        def resource_a(_):
            yield 'A'

        def resource_based_type_check(context, value):
            return context.resources.a == value

        CustomType = DagsterType(
            name='NeedsA',
            type_check_fn=resource_based_type_check,
            required_resource_keys={'a'} if should_require_resources else None,
        )

        @solid(output_defs=[OutputDefinition(CustomType, 'custom_type')])
        def custom_type_solid(_):
            return 'A'

        @pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a})])
        def type_check_pipeline():
            custom_type_solid()

        return type_check_pipeline

    under_required_pipeline = define_type_check_pipeline(should_require_resources=False)
    with pytest.raises(DagsterUnknownResourceError):
        execute_pipeline(under_required_pipeline)

    sufficiently_required_pipeline = define_type_check_pipeline(should_require_resources=True)
    assert execute_pipeline(sufficiently_required_pipeline).success
