# Copyright 2022 The Kubeflow Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Contains tests for kfp.dsl.placeholders."""
import os
import tempfile
from typing import Any, List

from absl.testing import parameterized
from kfp import compiler
from kfp import dsl
from kfp.dsl import Artifact
from kfp.dsl import Dataset
from kfp.dsl import Input
from kfp.dsl import Output
from kfp.dsl import placeholders


class TestInputValuePlaceholder(parameterized.TestCase):

    def test(self):
        self.assertEqual(
            placeholders.InputValuePlaceholder('input1')._to_string(),
            "{{$.inputs.parameters['input1']}}")


class TestInputPathPlaceholder(parameterized.TestCase):

    def test(self):
        self.assertEqual(
            placeholders.InputPathPlaceholder('input1')._to_string(),
            "{{$.inputs.artifacts['input1'].path}}")


class TestInputUriPlaceholder(parameterized.TestCase):

    def test(self):
        self.assertEqual(
            placeholders.InputUriPlaceholder('input1')._to_string(),
            "{{$.inputs.artifacts['input1'].uri}}")


class TestInputMetadataPlaceholder(parameterized.TestCase):

    def test(self):
        self.assertEqual(
            placeholders.InputMetadataPlaceholder('input1')._to_string(),
            "{{$.inputs.artifacts['input1'].metadata}}")

    def test_access_top_level_metadata_key(self):

        @dsl.container_component
        def echo(d: Input[Dataset]):
            return dsl.ContainerSpec(
                image='alpine', command=['echo', d.metadata['key']])

        self.assertEqual(
            echo.pipeline_spec.deployment_spec['executors']['exec-echo']
            ['container']['command'][1],
            "{{$.inputs.artifacts['d'].metadata['key']}}")


class TestOutputPathPlaceholder(parameterized.TestCase):

    def test(self):
        self.assertEqual(
            placeholders.OutputPathPlaceholder('output1')._to_string(),
            "{{$.outputs.artifacts['output1'].path}}")


class TestOutputParameterPlaceholder(parameterized.TestCase):

    def test(self):
        self.assertEqual(
            placeholders.OutputParameterPlaceholder('output1')._to_string(),
            "{{$.outputs.parameters['output1'].output_file}}")


class TestOutputUriPlaceholder(parameterized.TestCase):

    def test(self):
        self.assertEqual(
            placeholders.OutputUriPlaceholder('output1')._to_string(),
            "{{$.outputs.artifacts['output1'].uri}}")


class TestOutputMetadataPlaceholder(parameterized.TestCase):

    def test(self):
        self.assertEqual(
            placeholders.OutputMetadataPlaceholder('output1')._to_string(),
            "{{$.outputs.artifacts['output1'].metadata}}")

    def test_access_top_level_metadata_key(self):

        @dsl.container_component
        def echo(d: Output[Dataset]):
            return dsl.ContainerSpec(
                image='alpine', command=['echo', d.metadata['key']])

        self.assertEqual(
            echo.pipeline_spec.deployment_spec['executors']['exec-echo']
            ['container']['command'][1],
            "{{$.outputs.artifacts['d'].metadata['key']}}")


class TestIfPresentPlaceholder(parameterized.TestCase):

    @parameterized.parameters([
        (placeholders.IfPresentPlaceholder(
            input_name='input1', then=['then'], else_=['something']),
         '{"IfPresent": {"InputName": "input1", "Then": ["then"], "Else": ["something"]}}'
        ),
        (placeholders.IfPresentPlaceholder(
            input_name='input1', then='then', else_='something'),
         '{"IfPresent": {"InputName": "input1", "Then": "then", "Else": "something"}}'
        ),
        (placeholders.IfPresentPlaceholder(
            input_name='input1', then='then', else_=['something']),
         '{"IfPresent": {"InputName": "input1", "Then": "then", "Else": ["something"]}}'
        ),
        (placeholders.IfPresentPlaceholder(
            input_name='input1', then=['then'], else_='something'),
         '{"IfPresent": {"InputName": "input1", "Then": ["then"], "Else": "something"}}'
        ),
    ])
    def test_strings_and_lists(
            self, placeholder_obj: placeholders.IfPresentPlaceholder,
            placeholder: str):
        self.assertEqual(placeholder_obj._to_string(), placeholder)

    @parameterized.parameters([
        (placeholders.IfPresentPlaceholder(
            input_name='input1',
            then=[
                '--flag',
                placeholders.OutputUriPlaceholder(output_name='output1')
            ],
            else_=[
                '--flag',
                placeholders.OutputMetadataPlaceholder(output_name='output1')
            ]),
         """{"IfPresent": {"InputName": "input1", "Then": ["--flag", "{{$.outputs.artifacts['output1'].uri}}"], "Else": ["--flag", "{{$.outputs.artifacts['output1'].metadata}}"]}}"""
        ),
        (placeholders.IfPresentPlaceholder(
            input_name='input1',
            then=placeholders.InputPathPlaceholder(input_name='input2'),
            else_=placeholders.InputValuePlaceholder(input_name='input2')),
         """{"IfPresent": {"InputName": "input1", "Then": "{{$.inputs.artifacts['input2'].path}}", "Else": "{{$.inputs.parameters['input2']}}"}}"""
        ),
    ])
    def test_with_primitive_placeholders(
            self, placeholder_obj: placeholders.IfPresentPlaceholder,
            placeholder: str):
        self.assertEqual(placeholder_obj._to_string(), placeholder)

    def test_if_present_with_single_element_simple_can_be_compiled(self):

        @dsl.container_component
        def container_component(a: str):
            return dsl.ContainerSpec(
                image='alpine',
                command=[
                    placeholders.IfPresentPlaceholder(
                        input_name='a', then='b', else_='c')
                ])

        with tempfile.TemporaryDirectory() as tempdir:
            output_yaml = os.path.join(tempdir, 'component.yaml')
            compiler.Compiler().compile(
                pipeline_func=container_component, package_path=output_yaml)

    def test_if_present_with_single_element_parameter_reference_can_be_compiled(
            self):

        @dsl.container_component
        def container_component(a: str):
            return dsl.ContainerSpec(
                image='alpine',
                command=[
                    placeholders.IfPresentPlaceholder(
                        input_name='a', then=a, else_='c')
                ])

        with tempfile.TemporaryDirectory() as tempdir:
            output_yaml = os.path.join(tempdir, 'component.yaml')
            compiler.Compiler().compile(
                pipeline_func=container_component, package_path=output_yaml)

    def test_if_present_with_single_element_artifact_reference_can_be_compiled(
            self):

        @dsl.container_component
        def container_component(a: dsl.Input[dsl.Artifact]):
            return dsl.ContainerSpec(
                image='alpine',
                command=[
                    placeholders.IfPresentPlaceholder(
                        input_name='a', then=a.path, else_='c')
                ])

        with tempfile.TemporaryDirectory() as tempdir:
            output_yaml = os.path.join(tempdir, 'component.yaml')
            compiler.Compiler().compile(
                pipeline_func=container_component, package_path=output_yaml)


class TestConcatPlaceholder(parameterized.TestCase):

    @parameterized.parameters([
        (placeholders.ConcatPlaceholder(['a']), '{"Concat": ["a"]}'),
        (placeholders.ConcatPlaceholder(['a', 'b']), '{"Concat": ["a", "b"]}'),
    ])
    def test_strings(self, placeholder_obj: placeholders.ConcatPlaceholder,
                     placeholder_string: str):
        self.assertEqual(placeholder_obj._to_string(), placeholder_string)

    @parameterized.parameters([
        (placeholders.ConcatPlaceholder([
            'a', placeholders.InputPathPlaceholder(input_name='input2')
        ]), """{"Concat": ["a", "{{$.inputs.artifacts['input2'].path}}"]}"""),
        (placeholders.ConcatPlaceholder([
            placeholders.InputValuePlaceholder(input_name='input2'), 'b'
        ]), """{"Concat": ["{{$.inputs.parameters['input2']}}", "b"]}"""),
    ])
    def test_primitive_placeholders(
            self, placeholder_obj: placeholders.ConcatPlaceholder,
            placeholder_string: str):
        self.assertEqual(placeholder_obj._to_string(), placeholder_string)


class TestContainerPlaceholdersTogether(parameterized.TestCase):

    @parameterized.parameters([
        (placeholders.ConcatPlaceholder([
            'a', placeholders.ConcatPlaceholder(['b', 'c'])
        ]), '{"Concat": ["a", {"Concat": ["b", "c"]}]}'),
        (placeholders.ConcatPlaceholder([
            'a',
            placeholders.ConcatPlaceholder([
                'b',
                placeholders.ConcatPlaceholder([
                    'c',
                    placeholders.InputValuePlaceholder(input_name='input2')
                ])
            ])
        ]),
         """{"Concat": ["a", {"Concat": ["b", {"Concat": ["c", "{{$.inputs.parameters['input2']}}"]}]}]}"""
        ),
        (placeholders.ConcatPlaceholder([
            'a',
            placeholders.ConcatPlaceholder([
                'b',
                placeholders.IfPresentPlaceholder(
                    input_name='output1', then='then', else_='something')
            ])
        ]),
         '{"Concat": ["a", {"Concat": ["b", {"IfPresent": {"InputName": "output1", "Then": "then", "Else": "something"}}]}]}'
        ),
        (placeholders.ConcatPlaceholder([
            'a',
            placeholders.IfPresentPlaceholder(
                input_name='output1',
                then=placeholders.ConcatPlaceholder([
                    '--',
                    'flag',
                    placeholders.InputPathPlaceholder(input_name='input2'),
                ]),
                else_='b'),
            'c',
        ]),
         """{"Concat": ["a", {"IfPresent": {"InputName": "output1", "Then": {"Concat": ["--", "flag", "{{$.inputs.artifacts['input2'].path}}"]}, "Else": "b"}}, "c"]}"""
        ),
        (placeholders.ConcatPlaceholder([
            'a',
            placeholders.IfPresentPlaceholder(
                input_name='output1',
                then=placeholders.ConcatPlaceholder(['--', 'flag']),
                else_=placeholders.InputPathPlaceholder(input_name='input2')),
            'c',
        ]),
         """{"Concat": ["a", {"IfPresent": {"InputName": "output1", "Then": {"Concat": ["--", "flag"]}, "Else": "{{$.inputs.artifacts['input2'].path}}"}}, "c"]}"""
        ),
    ])
    def test_valid(self, placeholder_obj: placeholders.IfPresentPlaceholder,
                   placeholder: str):
        self.assertEqual(placeholder_obj._to_string(), placeholder)

    def test_only_single_element_ifpresent_inside_concat_outer(self):
        with self.assertRaisesRegex(
                ValueError,
                r'Please use a single element for `then` and `else_` only\.'):
            placeholders.ConcatPlaceholder([
                'b',
                placeholders.IfPresentPlaceholder(
                    input_name='output1', then=['then'], else_=['something'])
            ])

    def test_only_single_element_ifpresent_inside_concat_recursive(self):
        with self.assertRaisesRegex(
                ValueError,
                r'Please use a single element for `then` and `else_` only\.'):
            placeholders.ConcatPlaceholder([
                'a',
                placeholders.ConcatPlaceholder([
                    'b',
                    placeholders.IfPresentPlaceholder(
                        input_name='output1',
                        then=['then'],
                        else_=['something'])
                ])
            ])

        with self.assertRaisesRegex(
                ValueError,
                r'Please use a single element for `then` and `else_` only\.'):
            placeholders.ConcatPlaceholder([
                'a',
                placeholders.ConcatPlaceholder([
                    'b',
                    placeholders.IfPresentPlaceholder(
                        input_name='output1',
                        then=placeholders.ConcatPlaceholder([
                            placeholders.IfPresentPlaceholder(
                                input_name='a', then=['b'])
                        ]),
                        else_='something')
                ])
            ])

    def test_only_single_element_in_nested_ifpresent_inside_concat(self):
        with self.assertRaisesRegex(
                ValueError,
                r'Please use a single element for `then` and `else_` only\.'):
            dsl.ConcatPlaceholder([
                'my-prefix-',
                dsl.IfPresentPlaceholder(
                    input_name='input1',
                    then=[
                        dsl.IfPresentPlaceholder(
                            input_name='input1',
                            then=dsl.ConcatPlaceholder(['infix-', 'value']))
                    ])
            ])

    def test_recursive_nested_placeholder_validation_does_not_exit_when_first_valid_then_is_found(
            self):
        with self.assertRaisesRegex(
                ValueError,
                r'Please use a single element for `then` and `else_` only\.'):
            dsl.ConcatPlaceholder([
                'my-prefix-',
                dsl.IfPresentPlaceholder(
                    input_name='input1',
                    then=dsl.IfPresentPlaceholder(
                        input_name='input1',
                        then=[dsl.ConcatPlaceholder(['infix-', 'value'])]))
            ])

    def test_only_single_element_in_nested_ifpresent_inside_concat_with_outer_ifpresent(
            self):
        with self.assertRaisesRegex(
                ValueError,
                r'Please use a single element for `then` and `else_` only\.'):
            dsl.IfPresentPlaceholder(
                input_name='input_1',
                then=dsl.ConcatPlaceholder([
                    'my-prefix-',
                    dsl.IfPresentPlaceholder(
                        input_name='input1',
                        then=dsl.IfPresentPlaceholder(
                            input_name='input1',
                            then=[dsl.ConcatPlaceholder(['infix-', 'value'])]))
                ]))

    def test_valid_then_but_invalid_else(self):
        with self.assertRaisesRegex(
                ValueError,
                r'Please use a single element for `then` and `else_` only\.'):
            dsl.ConcatPlaceholder([
                'my-prefix-',
                dsl.IfPresentPlaceholder(
                    input_name='input1',
                    then=dsl.IfPresentPlaceholder(
                        input_name='input1',
                        then='single-element',
                        else_=['one', 'two']))
            ])


class TestListOfArtifactsInContainerComponentPlaceholders(
        parameterized.TestCase):

    def test_compile_component1(self):

        @dsl.container_component
        def comp(input_list: Input[List[Artifact]]):
            return dsl.ContainerSpec(
                image='alpine', command=[input_list], args=[input_list])

        self.assertEqual(
            comp.pipeline_spec.deployment_spec['executors']['exec-comp']
            ['container']['command'][0], "{{$.inputs.artifacts['input_list']}}")
        self.assertEqual(
            comp.pipeline_spec.deployment_spec['executors']['exec-comp']
            ['container']['args'][0], "{{$.inputs.artifacts['input_list']}}")

    def test_compile_component2(self):

        @dsl.container_component
        def comp(new_name: Input[List[Dataset]]):
            return dsl.ContainerSpec(
                image='alpine', command=[new_name], args=[new_name])

        self.assertEqual(
            comp.pipeline_spec.deployment_spec['executors']['exec-comp']
            ['container']['command'][0], "{{$.inputs.artifacts['new_name']}}")
        self.assertEqual(
            comp.pipeline_spec.deployment_spec['executors']['exec-comp']
            ['container']['args'][0], "{{$.inputs.artifacts['new_name']}}")

    def test_cannot_access_name(self):
        with self.assertRaisesRegex(AttributeError,
                                    'Cannot access an attribute'):

            @dsl.container_component
            def comp(new_name: Input[List[Dataset]]):
                return dsl.ContainerSpec(
                    image='alpine', command=[new_name.name])

    def test_cannot_access_uri(self):
        with self.assertRaisesRegex(AttributeError,
                                    'Cannot access an attribute'):

            @dsl.container_component
            def comp(new_name: Input[List[Dataset]]):
                return dsl.ContainerSpec(image='alpine', command=[new_name.uri])

    def test_cannot_access_metadata(self):
        with self.assertRaisesRegex(AttributeError,
                                    'Cannot access an attribute'):

            @dsl.container_component
            def comp(new_name: Input[List[Dataset]]):
                return dsl.ContainerSpec(
                    image='alpine', command=[new_name.metadata])

    def test_cannot_access_path(self):
        with self.assertRaisesRegex(AttributeError,
                                    'Cannot access an attribute'):

            @dsl.container_component
            def comp(new_name: Input[List[Dataset]]):
                return dsl.ContainerSpec(
                    image='alpine', command=[new_name.path])

    def test_cannot_access_individual_artifact(self):
        with self.assertRaisesRegex(KeyError,
                                    'Cannot access individual artifacts'):

            @dsl.container_component
            def comp(new_name: Input[List[Dataset]]):
                return dsl.ContainerSpec(image='alpine', command=[new_name[0]])


class TestConvertCommandLineElementToStringOrStruct(parameterized.TestCase):

    @parameterized.parameters(['a', 'word', 1])
    def test_pass_through(self, val: Any):
        self.assertEqual(
            placeholders.convert_command_line_element_to_string_or_struct(val),
            val)

    @parameterized.parameters([
        (placeholders.InputValuePlaceholder('input1'),
         """{{$.inputs.parameters['input1']}}"""),
        (placeholders.OutputPathPlaceholder('output1'),
         """{{$.outputs.artifacts['output1'].path}}"""),
    ])
    def test_primitive_placeholder(self, placeholder: Any, expected: str):
        self.assertEqual(
            placeholders.convert_command_line_element_to_string_or_struct(
                placeholder), expected)


class OtherPlaceholderTests(parameterized.TestCase):

    def test_primitive_placeholder_can_be_used_in_fstring1(self):

        @dsl.container_component
        def echo_bool(boolean: bool = True):
            return dsl.ContainerSpec(
                image='alpine', command=['sh', '-c', f'echo {boolean}'])

        self.assertEqual(
            echo_bool.component_spec.implementation.container.command,
            ['sh', '-c', "echo {{$.inputs.parameters['boolean']}}"])

    def test_primitive_placeholder_can_be_used_in_fstring2(self):

        @dsl.container_component
        def container_with_placeholder_in_fstring(
            output_artifact: Output[Artifact],
            text1: str,
        ):
            return dsl.ContainerSpec(
                image='python:3.9',
                command=[
                    'my_program',
                    f'prefix-{text1}',
                    f'{output_artifact.uri}/0',
                ])

        self.assertEqual(
            container_with_placeholder_in_fstring.component_spec.implementation
            .container.command, [
                'my_program',
                "prefix-{{$.inputs.parameters['text1']}}",
                "{{$.outputs.artifacts['output_artifact'].uri}}/0",
            ])

    def test_cannot_use_concat_placeholder_in_f_string(self):

        with self.assertRaisesRegex(
                ValueError, 'Cannot use ConcatPlaceholder in an f-string.'):

            @dsl.container_component
            def container_with_placeholder_in_fstring(
                text1: str,
                text2: str,
            ):
                return dsl.ContainerSpec(
                    image='python:3.9',
                    command=[
                        'my_program',
                        f'another-prefix-{dsl.ConcatPlaceholder([text1, text2])}',
                    ])

    def test_cannot_use_ifpresent_placeholder_in_f_string(self):

        with self.assertRaisesRegex(
                ValueError, 'Cannot use IfPresentPlaceholder in an f-string.'):

            @dsl.container_component
            def container_with_placeholder_in_fstring(
                text1: str,
                text2: str,
            ):
                return dsl.ContainerSpec(
                    image='python:3.9',
                    command=[
                        'echo',
                        f"another-prefix-{dsl.IfPresentPlaceholder(input_name='text1', then=['val'])}",
                    ])
