from pathlib import Path

import pytest

from azure.ai.ml import load_component, Input
from azure.ai.ml.dsl import pipeline
from azure.ai.ml.entities._job.pipeline._io import PipelineInput
from azure.ai.ml.entities._job.pipeline._io.base import _resolve_builders_2_data_bindings
from azure.ai.ml.exceptions import UserErrorException
from test_utilities.utils import omit_with_wildcard

from .._util import _DSL_TIMEOUT_SECOND

tests_root_dir = Path(__file__).parent.parent.parent
components_dir = tests_root_dir / "test_configs/components/"
common_omit_fields = [
    "jobs.*.componentId",
    "jobs.*._source",
    "jobs.*.properties",
]


@pytest.mark.usefixtures("enable_pipeline_private_preview_features")
@pytest.mark.timeout(_DSL_TIMEOUT_SECOND)
@pytest.mark.unittest
@pytest.mark.pipeline_test
class TestInputOutputBuilder:
    def test_nested_input_output_builder(self):
        input1 = PipelineInput(name="input1", owner="pipeline", meta=None)
        input2 = PipelineInput(name="input2", owner="pipeline", meta=None)
        input3 = PipelineInput(name="input3", owner="pipeline", meta=None)
        test_data = {
            "data1": input1,
            "data2": {
                "1": input2,
            },
            "data3": [input3],
            "data4": [{"1": input1}, input2, [input3]],
            "data5": {"1": [input1], "2": {"1": input2}, "3": input3},
        }
        assert _resolve_builders_2_data_bindings(test_data) == {
            "data1": "${{parent.inputs.input1}}",
            "data2": {"1": "${{parent.inputs.input2}}"},
            "data3": ["${{parent.inputs.input3}}"],
            "data4": [{"1": "${{parent.inputs.input1}}"}, "${{parent.inputs.input2}}", ["${{parent.inputs.input3}}"]],
            "data5": {
                "1": ["${{parent.inputs.input1}}"],
                "2": {"1": "${{parent.inputs.input2}}"},
                "3": "${{parent.inputs.input3}}",
            },
        }

    def test_pipeline_input_result(self):
        component_yaml = components_dir / "helloworld_component.yml"
        component_func = load_component(source=component_yaml)

        @pipeline
        def my_pipeline(job_in_number, job_in_path):
            assert isinstance(job_in_number, PipelineInput)
            assert isinstance(job_in_path, PipelineInput)
            assert isinstance(job_in_number.result(), int)
            assert isinstance(job_in_path.result(), Input)
            node1 = component_func(component_in_number=job_in_number, component_in_path=job_in_path)
            # calling result() will convert pipeline input to actual value
            node2 = component_func(component_in_number=job_in_number.result(), component_in_path=job_in_path.result())
            return {
                "output1": node1.outputs.component_out_path,
                "output2": node2.outputs.component_out_path
            }

        pipeline_job1 = my_pipeline(
            job_in_number=1, job_in_path=Input(path="fake_path1")
        )

        rest_pipeline_job = omit_with_wildcard(pipeline_job1._to_rest_object().properties.as_dict(),
                                               *common_omit_fields)
        expected_pipeline_job1 = {
            'node1': {
                'inputs': {'component_in_number': {'job_input_type': 'literal',
                                                   'value': '${{parent.inputs.job_in_number}}'},
                           'component_in_path': {'job_input_type': 'literal',
                                                 'value': '${{parent.inputs.job_in_path}}'}},
                'name': 'node1',
                'outputs': {'component_out_path': {'type': 'literal',
                                                   'value': '${{parent.outputs.output1}}'}},
                'type': 'command'},
            'node2': {
                'inputs': {'component_in_number': {'job_input_type': 'literal',
                                                   'value': '1'},
                           'component_in_path': {'job_input_type': 'uri_folder',
                                                 'uri': 'fake_path1'}},
                'name': 'node2',
                'outputs': {'component_out_path': {'type': 'literal',
                                                   'value': '${{parent.outputs.output2}}'}},
                'type': 'command'}
        }
        assert rest_pipeline_job["jobs"] == expected_pipeline_job1

        pipeline_job2 = my_pipeline(
            job_in_number=2, job_in_path=Input(path="fake_path2")
        )

        rest_pipeline_job = omit_with_wildcard(pipeline_job2._to_rest_object().properties.as_dict(),
                                               *common_omit_fields)

        expected_pipeline_job2 = {
            'node1': {'inputs': {'component_in_number': {'job_input_type': 'literal',
                                                         'value': '${{parent.inputs.job_in_number}}'},
                                 'component_in_path': {'job_input_type': 'literal',
                                                       'value': '${{parent.inputs.job_in_path}}'}},
                      'name': 'node1',
                      'outputs': {'component_out_path': {'type': 'literal',
                                                         'value': '${{parent.outputs.output1}}'}},
                      'type': 'command'},
            'node2': {'inputs': {'component_in_number': {'job_input_type': 'literal',
                                                         'value': '2'},
                                 'component_in_path': {'job_input_type': 'uri_folder',
                                                       'uri': 'fake_path2'}},
                      'name': 'node2',
                      'outputs': {'component_out_path': {'type': 'literal',
                                                         'value': '${{parent.outputs.output2}}'}},
                      'type': 'command'}
        }
        assert rest_pipeline_job["jobs"] == expected_pipeline_job2

        # calling pipeline func again won't affect existing pipeline job
        pipeline_job1.jobs["node2"].inputs["component_in_number"] == 1
        pipeline_job1.jobs["node2"].inputs["component_in_path"].path == "fake_path1"

        rest_pipeline_job = omit_with_wildcard(pipeline_job1._to_rest_object().properties.as_dict(),
                                               *common_omit_fields)
        assert rest_pipeline_job["jobs"] == expected_pipeline_job1

    def test_pipeline_input_result_multiple_levle(self):
        component_yaml = components_dir / "helloworld_component.yml"
        component_func = load_component(source=component_yaml)

        @pipeline
        def my_pipeline_level_1(job_in_number, job_in_path):
            assert isinstance(job_in_number, PipelineInput)
            assert isinstance(job_in_path, PipelineInput)
            # Note: call result will get actual value
            assert isinstance(job_in_number.result(), int)
            assert isinstance(job_in_path.result(), Input)
            component_func(component_in_number=job_in_number, component_in_path=job_in_path)

        @pipeline
        def my_pipeline_level_2(job_in_number, job_in_path):
            assert isinstance(job_in_number, PipelineInput)
            assert isinstance(job_in_path, PipelineInput)
            assert isinstance(job_in_number.result(), int)
            assert isinstance(job_in_path.result(), Input)
            my_pipeline_level_1(job_in_number=job_in_number, job_in_path=job_in_path)
            component_func(component_in_number=job_in_number, component_in_path=job_in_path)

        pipeline_job2 = my_pipeline_level_2(
            job_in_number=2, job_in_path=Input(path="fake_path2")
        )

        rest_pipeline_job = omit_with_wildcard(pipeline_job2._to_rest_object().properties.as_dict(),
                                               *common_omit_fields)

        expected_pipeline_job = {
            'microsoftsamples_command_component_basic': {
                'inputs': {'component_in_number': {'job_input_type': 'literal',
                                                   'value': '${{parent.inputs.job_in_number}}'},
                           'component_in_path': {'job_input_type': 'literal',
                                                 'value': '${{parent.inputs.job_in_path}}'}},
                'name': 'microsoftsamples_command_component_basic',
                'type': 'command'},
            'my_pipeline_level_1': {
                'inputs': {'job_in_number': {'job_input_type': 'literal',
                                             'value': '${{parent.inputs.job_in_number}}'},
                           'job_in_path': {'job_input_type': 'literal',
                                           'value': '${{parent.inputs.job_in_path}}'}},
                'name': 'my_pipeline_level_1',
                'type': 'pipeline'}
        }
        assert rest_pipeline_job["jobs"] == expected_pipeline_job

    def test_pipeline_expression_bool_test(self) -> None:
        # non-pipeline scenario, bool test will return True
        input1 = PipelineInput(name="input1", owner="pipeline", meta=None)
        if input1:
            pass
        else:
            assert False, "bool test for PipelineInput in non-pipeline scenario should always return True."

        # pipeline scenario, should raise UserErrorException
        @pipeline
        def pipeline_func(int_param: int):
            if int_param:
                print("should not enter this line.")

        with pytest.raises(UserErrorException) as e:
            pipeline_func(int_param=1)
        assert str(e.value) == (
            "Type <class 'azure.ai.ml.entities._job.pipeline._io.base.PipelineInput'> "
            "is not supported for operation bool()."
        )
