from dagster import Field, Int, Noneable, PipelineDefinition, String, solid
from dagster.config.field import resolve_to_config_type
from dagster.config.type_printer import print_type_to_string


def assert_inner_types(parent_type, *dagster_types):
    assert set(
        list(map(lambda t: t.key, resolve_to_config_type(parent_type).recursive_config_types))
    ) == set(map(lambda x: x.key, map(resolve_to_config_type, dagster_types)))


def test_basic_type_print():
    assert print_type_to_string(Int) == 'Int'
    assert_inner_types(Int)


def test_basic_list_type_print():
    assert print_type_to_string([int]) == '[Int]'
    assert_inner_types([int], Int)


def test_double_list_type_print():
    assert print_type_to_string([[int]]) == '[[Int]]'
    int_list = [int]
    list_int_list = [int_list]
    assert_inner_types(list_int_list, Int, int_list)


def test_basic_nullable_type_print():
    assert print_type_to_string(Noneable(int)) == 'Int?'
    nullable_int = Noneable(int)
    assert_inner_types(nullable_int, Int)


def test_nullable_list_combos():
    assert print_type_to_string([int]) == '[Int]'
    assert print_type_to_string(Noneable([int])) == '[Int]?'
    assert print_type_to_string([Noneable(int)]) == '[Int?]'
    assert print_type_to_string(Noneable([Noneable(int)])) == '[Int?]?'


def test_basic_dict():
    output = print_type_to_string({'int_field': int})

    expected = '''{
  int_field: Int
}'''

    assert output == expected


def test_two_field_dicts():
    two_field_dict = {'int_field': int, 'string_field': str}
    assert_inner_types(two_field_dict, Int, String)

    output = print_type_to_string(two_field_dict)

    expected = '''{
  int_field: Int
  string_field: String
}'''

    assert output == expected


def test_two_field_dicts_same_type():
    two_field_dict = {'int_field1': int, 'int_field2': int}
    assert_inner_types(two_field_dict, Int)

    output = print_type_to_string(two_field_dict)

    expected = '''{
  int_field1: Int
  int_field2: Int
}'''

    assert output == expected


def test_optional_field():
    output = print_type_to_string({'int_field': Field(int, is_required=False)})

    expected = '''{
  int_field?: Int
}'''

    assert output == expected


def test_single_level_dict_lists_and_nullable():
    output = print_type_to_string(
        {
            'nullable_int_field': Noneable(int),
            'optional_int_field': Field(int, is_required=False),
            'string_list_field': [str],
        }
    )

    expected = '''{
  nullable_int_field: Int?
  optional_int_field?: Int
  string_list_field: [String]
}'''

    assert output == expected


def test_nested_dict():
    nested_type = {'int_field': int}
    outer_type = {'nested': nested_type}
    output = print_type_to_string(outer_type)

    assert_inner_types(outer_type, Int, nested_type)

    expected = '''{
  nested: {
    int_field: Int
  }
}'''

    assert output == expected


def test_test_type_pipeline_construction():
    assert define_test_type_pipeline()


def define_solid_for_test_type(name, config):
    @solid(name=name, config=config, input_defs=[], output_defs=[])
    def a_solid(_):
        return None

    return a_solid


# launch in dagit with this command:
# dagit -f test_type_printer.py -n define_test_type_pipeline
def define_test_type_pipeline():
    return PipelineDefinition(
        name='test_type_pipeline',
        solid_defs=[
            define_solid_for_test_type('int_config', int),
            define_solid_for_test_type('list_of_int_config', [int]),
            define_solid_for_test_type('nullable_list_of_int_config', Noneable([int])),
            define_solid_for_test_type('list_of_nullable_int_config', [Noneable(int)]),
            define_solid_for_test_type(
                'nullable_list_of_nullable_int_config', Noneable([Noneable(int)])
            ),
            define_solid_for_test_type('simple_dict', {'int_field': int, 'string_field': str}),
            define_solid_for_test_type(
                'dict_with_optional_field',
                {
                    'nullable_int_field': Noneable(int),
                    'optional_int_field': Field(int, is_required=False),
                    'string_list_field': [str],
                },
            ),
            define_solid_for_test_type('nested_dict', {'nested': {'int_field': int}}),
        ],
    )
