# Copyright 2020-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.
"""Utilities for component I/O type mapping."""

from distutils import util
import inspect
import json
from typing import Any, Callable, Dict, Optional, Type, Union
import warnings

import kfp
from kfp.dsl import structures
from kfp.dsl import task_final_status
from kfp.dsl.types import artifact_types
from kfp.dsl.types import type_annotations
from kfp.pipeline_spec import pipeline_spec_pb2

DEFAULT_ARTIFACT_SCHEMA_VERSION = '0.0.1'
PARAMETER_TYPES = Union[str, int, float, bool, dict, list]

# ComponentSpec I/O types to DSL ontology artifact classes mapping.
_ARTIFACT_CLASSES_MAPPING = {
    'artifact': artifact_types.Artifact,
    'model': artifact_types.Model,
    'dataset': artifact_types.Dataset,
    'metrics': artifact_types.Metrics,
    'classificationmetrics': artifact_types.ClassificationMetrics,
    'slicedclassificationmetrics': artifact_types.SlicedClassificationMetrics,
    'html': artifact_types.HTML,
    'markdown': artifact_types.Markdown,
}

_GOOGLE_TYPES_PATTERN = r'^google.[A-Za-z]+$'
_GOOGLE_TYPES_VERSION = DEFAULT_ARTIFACT_SCHEMA_VERSION

# ComponentSpec I/O types to (IR) PipelineTaskSpec I/O types mapping.
# The keys are normalized (lowercased). These are types viewed as Parameters.
# The values are the corresponding IR parameter primitive types.
_PARAMETER_TYPES_MAPPING = {
    'integer': pipeline_spec_pb2.ParameterType.NUMBER_INTEGER,
    'int': pipeline_spec_pb2.ParameterType.NUMBER_INTEGER,
    'double': pipeline_spec_pb2.ParameterType.NUMBER_DOUBLE,
    'float': pipeline_spec_pb2.ParameterType.NUMBER_DOUBLE,
    'string': pipeline_spec_pb2.ParameterType.STRING,
    'str': pipeline_spec_pb2.ParameterType.STRING,
    'text': pipeline_spec_pb2.ParameterType.STRING,
    'bool': pipeline_spec_pb2.ParameterType.BOOLEAN,
    'boolean': pipeline_spec_pb2.ParameterType.BOOLEAN,
    'dict': pipeline_spec_pb2.ParameterType.STRUCT,
    'list': pipeline_spec_pb2.ParameterType.LIST,
    'jsonobject': pipeline_spec_pb2.ParameterType.STRUCT,
    'jsonarray': pipeline_spec_pb2.ParameterType.LIST,
}


def bool_cast_fn(default: Union[str, bool]) -> bool:
    if isinstance(default, str):
        default = util.strtobool(default) == 1
    return default


def try_loading_json(default: str) -> Union[dict, list, str]:
    try:
        return json.loads(default)
    except:
        return default


_V1_DEFAULT_DESERIALIZER_MAPPING: Dict[str, Callable] = {
    'integer': int,
    'int': int,
    'double': float,
    'float': float,
    'string': str,
    'str': str,
    'text': str,
    'bool': bool_cast_fn,
    'boolean': bool_cast_fn,
    'dict': try_loading_json,
    'list': try_loading_json,
    'jsonobject': try_loading_json,
    'jsonarray': try_loading_json,
}


def deserialize_v1_component_yaml_default(type_: str, default: Any) -> Any:
    """Deserializes v1 default values to correct in-memory types.

    Typecasts for primitive types. Tries to load JSON for arrays and
    structs.
    """
    if default is None:
        return default
    if isinstance(type_, str):
        cast_fn = _V1_DEFAULT_DESERIALIZER_MAPPING.get(type_.lower(),
                                                       lambda x: x)
        return cast_fn(default)
    return default


# Mapping primitive types to their IR message field names.
# This is used in constructing condition strings.
_PARAMETER_TYPES_VALUE_REFERENCE_MAPPING = {
    pipeline_spec_pb2.ParameterType.NUMBER_INTEGER: 'number_value',
    pipeline_spec_pb2.ParameterType.NUMBER_DOUBLE: 'number_value',
    pipeline_spec_pb2.ParameterType.STRING: 'string_value',
    pipeline_spec_pb2.ParameterType.BOOLEAN: 'bool_value',
    pipeline_spec_pb2.ParameterType.STRUCT: 'struct_value',
    pipeline_spec_pb2.ParameterType.LIST: 'list_value',
}


def is_task_final_status_type(type_name: Optional[Union[str, dict]]) -> bool:
    """Check if a ComponentSpec I/O type is PipelineTaskFinalStatus.

    Args:
      type_name: type name of the ComponentSpec I/O type.

    Returns:
      True if the type name is 'PipelineTaskFinalStatus'.
    """
    return isinstance(type_name, str) and (
        type_name == task_final_status.PipelineTaskFinalStatus.__name__)


def is_parameter_type(type_name: Optional[Union[str, dict]]) -> bool:
    """Check if a ComponentSpec I/O type is considered as a parameter type.

    Args:
      type_name: type name of the ComponentSpec I/O type.

    Returns:
      True if the type name maps to a parameter type else False.
    """
    if isinstance(type_name, str):
        type_name = type_annotations.get_short_type_name(type_name)
    elif isinstance(type_name, dict):
        type_name = list(type_name.keys())[0]
    else:
        return False

    return type_name.lower(
    ) in _PARAMETER_TYPES_MAPPING or is_task_final_status_type(type_name)


def bundled_artifact_to_artifact_proto(
        bundled_artifact_str: str) -> pipeline_spec_pb2.ArtifactTypeSchema:
    """Gets the IR ArtifactTypeSchema proto for a bundled artifact in form
    `<namespace>.<Name>@x.x.x` (e.g., system.Artifact@0.0.1)."""
    bundled_artifact_str, schema_version = bundled_artifact_str.split('@')
    return pipeline_spec_pb2.ArtifactTypeSchema(
        schema_title=bundled_artifact_str,
        schema_version=schema_version,
    )


def get_parameter_type(
    param_type: Optional[Union[Type, str, dict]]
) -> pipeline_spec_pb2.ParameterType:
    """Get the IR I/O parameter type for the given ComponentSpec I/O type.

    Args:
      param_type: type of the ComponentSpec I/O type. Can be a primitive Python
        builtin type or a type name.

    Returns:
      The enum value of the mapped IR I/O primitive type.

    Raises:
      AttributeError: if type_name is not a string type.
    """
    # Special handling for PipelineTaskFinalStatus, treat it as Dict type.
    if is_task_final_status_type(param_type):
        param_type = 'dict'
    if type(param_type) == type:
        type_name = param_type.__name__
    elif isinstance(param_type, dict):
        type_name = list(param_type.keys())[0]
    else:
        type_name = type_annotations.get_short_type_name(str(param_type))
    return _PARAMETER_TYPES_MAPPING.get(type_name.lower())


def get_parameter_type_name(
        param_type: Optional[Union[Type, str, dict]]) -> str:
    """Gets the parameter type name."""
    return pipeline_spec_pb2.ParameterType.ParameterTypeEnum.Name(
        get_parameter_type(param_type))


def get_parameter_type_field_name(type_name: Optional[str]) -> Optional[str]:
    """Get the IR field name for the given primitive type.

    For example: 'str' -> 'string_value', 'double' -> 'double_value', etc.

    Args:
      type_name: type name of the ComponentSpec I/O primitive type.

    Returns:
      The IR value reference field name.

    Raises:
      AttributeError: if type_name is not a string type.
    """
    return _PARAMETER_TYPES_VALUE_REFERENCE_MAPPING.get(
        get_parameter_type(type_name))


class InconsistentTypeException(Exception):
    """InconsistencyTypeException is raised when two types are not
    consistent."""


class InconsistentTypeWarning(Warning):
    """InconsistentTypeWarning is issued when two types are not consistent."""


def _get_type_string_from_component_argument(
    argument_value: Union['pipeline_channel.PipelineChannel', str, bool, int,
                          float, dict, list]
) -> str:
    # avoid circular imports
    from kfp.dsl import pipeline_channel
    if isinstance(argument_value, pipeline_channel.PipelineChannel):
        return argument_value.channel_type

    # argument is a constant
    argument_type = type(argument_value)
    if argument_type in _TYPE_TO_TYPE_NAME:
        return _TYPE_TO_TYPE_NAME[argument_type]

    raise ValueError(
        f'Constant argument inputs must be one of type {list(_TYPE_TO_TYPE_NAME.values())} Got: {argument_value!r} of type {type(argument_value)!r}.'
    )


def verify_type_compatibility(
    given_value: Union['pipeline_channel.PipelineChannel', str, bool, int,
                       float, dict, list],
    expected_spec: Union[structures.InputSpec, structures.OutputSpec],
    error_message_prefix: str,
    checks_input: bool = True,
) -> bool:
    """Verifies the given argument type is compatible with the expected type.

    Args:
        given_value: The channel or constant provided as an argument.
        expected_spec: The InputSpec or OutputSpec that describes the expected type of given_value.
        error_message_prefix: The prefix for the error message.
        checks_input: True if checks an argument (given_value) against a component/pipeline input type (expected_spec). False if checks a component output (argument_value) against the pipeline output type (expected_spec).

    Returns:
        True if types are compatible, and False if otherwise.

    Raises:
        InconsistentTypeException if types are incompatible and TYPE_CHECK==True.
    """
    # extract and normalize types
    expected_type = expected_spec.type
    given_type = _get_type_string_from_component_argument(given_value)

    given_is_param = is_parameter_type(str(given_type))
    if given_is_param:
        given_type = get_parameter_type_name(given_type)
        given_is_artifact_list = False
    else:
        given_is_artifact_list = given_value.is_artifact_list

    expected_is_param = is_parameter_type(expected_type)
    if expected_is_param:
        expected_type = get_parameter_type_name(expected_type)
        expected_is_artifact_list = False
    else:
        expected_is_artifact_list = expected_spec.is_artifact_list

    # compare the normalized types
    if given_is_param != expected_is_param:
        types_are_compatible = False
    elif given_is_param and expected_is_param:
        types_are_compatible = check_parameter_type_compatibility(
            given_type, expected_type)
    else:
        types_are_compatible = check_artifact_type_compatibility(
            given_type=given_type,
            given_is_artifact_list=given_is_artifact_list,
            expected_type=expected_type,
            expected_is_artifact_list=expected_is_artifact_list)

    # maybe raise, maybe warn, return bool
    if not types_are_compatible:
        # update the types for lists of artifacts for error message
        given_type = f'List[{given_type}]' if given_is_artifact_list else given_type
        expected_type = f'List[{expected_type}]' if expected_is_artifact_list else expected_type
        if checks_input:
            error_message_suffix = f'Argument type {given_type!r} is incompatible with the input type {expected_type!r}'
        else:
            error_message_suffix = f'Output of type {given_type!r} cannot be surfaced as pipeline output type {expected_type!r}'
        error_text = error_message_prefix + error_message_suffix
        if kfp.TYPE_CHECK:
            raise InconsistentTypeException(error_text)
        else:
            warnings.warn(InconsistentTypeWarning(error_text))

    return types_are_compatible


def check_artifact_type_compatibility(given_type: str,
                                      given_is_artifact_list: bool,
                                      expected_type: str,
                                      expected_is_artifact_list: bool) -> bool:
    given_schema_title, given_schema_version = given_type.split('@')
    expected_schema_title, expected_schema_version = expected_type.split('@')
    same_list_of_artifacts_status = expected_is_artifact_list == given_is_artifact_list
    if not same_list_of_artifacts_status:
        return False
    elif artifact_types.Artifact.schema_title in {
            given_schema_title, expected_schema_title
    }:
        return True
    else:
        schema_title_compatible = given_schema_title == expected_schema_title
        schema_version_compatible = given_schema_version.split(
            '.')[0] == expected_schema_version.split('.')[0]

        return schema_title_compatible and schema_version_compatible


def check_parameter_type_compatibility(given_type: str,
                                       expected_type: str) -> bool:
    if isinstance(given_type, str) and isinstance(expected_type, str):
        return given_type == expected_type
    else:
        return check_v1_struct_parameter_type_compatibility(
            given_type, expected_type)


def check_v1_struct_parameter_type_compatibility(
    given_type: Union[str, dict],
    expected_type: Union[str, dict],
) -> bool:
    if isinstance(given_type, str):
        given_type = {given_type: {}}
    if isinstance(expected_type, str):
        expected_type = {expected_type: {}}
    return _check_dict_types(given_type, expected_type)


def _check_dict_types(
    given_type: dict,
    expected_type: dict,
) -> bool:
    given_type_name, _ = list(given_type.items())[0]
    expected_type_name, _ = list(expected_type.items())[0]
    if given_type_name == '' or expected_type_name == '':
        # If the type name is empty, it matches any types
        return True
    if given_type_name != expected_type_name:
        print('type name ' + str(given_type_name) +
              ' is different from expected: ' + str(expected_type_name))
        return False
    type_name = given_type_name
    for type_property in given_type[type_name]:
        if type_property not in expected_type[type_name]:
            print(type_name + ' has a property ' + str(type_property) +
                  ' that the latter does not.')
            return False
        if given_type[type_name][type_property] != expected_type[type_name][
                type_property]:
            print(type_name + ' has a property ' + str(type_property) +
                  ' with value: ' + str(given_type[type_name][type_property]) +
                  ' and ' + str(expected_type[type_name][type_property]))
            return False
    return True


_TYPE_TO_TYPE_NAME = {
    str: 'String',
    int: 'Integer',
    float: 'Float',
    bool: 'Boolean',
    list: 'List',
    dict: 'Dict',
}


def get_canonical_type_name_for_type(typ: Type) -> Optional[str]:
    """Find the canonical type name for a given type.

    Args:
        typ: The type to search for.

    Returns:
        The canonical name of the type found.
    """
    return _TYPE_TO_TYPE_NAME.get(typ, None)


class TypeCheckManager:
    """Context manager to set a type check mode within context, then restore
    mode to original value upon exiting the context."""

    def __init__(self, enable: bool) -> None:
        """TypeCheckManager constructor.

        Args:
            enable: Type check mode used within context.
        """
        self._enable = enable

    def __enter__(self) -> 'TypeCheckManager':
        """Set type check mode to self._enable.

        Returns:
            TypeCheckManager: Returns itself.
        """
        self._prev = kfp.TYPE_CHECK
        kfp.TYPE_CHECK = self._enable
        return self

    def __exit__(self, *unused_args) -> None:
        """Restore type check mode to its previous state."""
        kfp.TYPE_CHECK = self._prev


# for reading in IR back to in-memory data structures
IR_TYPE_TO_IN_MEMORY_SPEC_TYPE = {
    'STRING': 'String',
    'NUMBER_INTEGER': 'Integer',
    'NUMBER_DOUBLE': 'Float',
    'LIST': 'List',
    'STRUCT': 'Dict',
    'BOOLEAN': 'Boolean',
    'TASK_FINAL_STATUS': task_final_status.PipelineTaskFinalStatus.__name__,
}

IR_TYPE_TO_COMMENT_TYPE_STRING = {
    'STRING': str.__name__,
    'NUMBER_INTEGER': int.__name__,
    'NUMBER_DOUBLE': float.__name__,
    'LIST': list.__name__,
    'STRUCT': dict.__name__,
    'BOOLEAN': bool.__name__,
    'TASK_FINAL_STATUS': task_final_status.PipelineTaskFinalStatus.__name__,
}

IN_MEMORY_SPEC_TYPE_TO_IR_TYPE = {
    v: k for k, v in IR_TYPE_TO_IN_MEMORY_SPEC_TYPE.items()
}


def get_canonical_name_for_outer_generic(type_name: Any) -> str:
    """Maps a complex/nested type name back to a canonical type.

    E.g.
        get_canonical_name_for_outer_generic('typing.List[str]')
        'List'

        get_canonical_name_for_outer_generic('typing.Dict[typing.List[str], str]')
        'Dict'

    Args:
        type_name (Any): The type. Returns input if not a string.

    Returns:
        str: The canonical type.
    """
    if not isinstance(type_name, str):
        return type_name

    if type_name.startswith('typing.'):
        type_name = type_name.lstrip('typing.')

    if type_name.lower().startswith('list') or type_name.lower().startswith(
            'dict'):
        return type_name.split('[')[0]

    else:
        return type_name


def create_bundled_artifact_type(schema_title: str,
                                 schema_version: Optional[str] = None) -> str:
    if not isinstance(schema_title, str):
        raise ValueError
    return schema_title + '@' + (
        schema_version or DEFAULT_ARTIFACT_SCHEMA_VERSION)


def validate_schema_version(schema_version: str) -> None:
    split_schema_version = schema_version.split('.')
    if len(split_schema_version) != 3:
        raise TypeError(
            f'Artifact schema_version must use three-part semantic versioning. Got: {schema_version}'
        )


def validate_schema_title(schema_title: str) -> None:
    split_schema_title = schema_title.split('.')
    if len(split_schema_title) != 2:
        raise TypeError(
            f'Artifact schema_title must have both a namespace and a name, separated by a `.`. Got: {schema_title}'
        )
    namespace, _ = split_schema_title
    if namespace not in {'system', 'google'}:
        raise TypeError(
            f'Artifact schema_title must belong to `system` or `google` namespace. Got: {schema_title}'
        )


def validate_bundled_artifact_type(type_: str) -> None:
    split_type = type_.split('@')
    # two parts and neither are empty strings
    if len(split_type) != 2 or not all(split_type):
        raise TypeError(
            f'Artifacts must have both a schema_title and a schema_version, separated by `@`. Got: {type_}'
        )
    schema_title, schema_version = split_type
    validate_schema_title(schema_title)
    validate_schema_version(schema_version)


def _annotation_to_type_struct(annotation):
    if not annotation or annotation == inspect.Parameter.empty:
        return None
    if hasattr(annotation, 'to_dict'):
        annotation = annotation.to_dict()
    if isinstance(annotation, dict):
        return annotation
    if isinstance(annotation, type):
        type_struct = get_canonical_type_name_for_type(annotation)
        if type_struct:
            return type_struct
        elif type_annotations.is_artifact_class(annotation):
            schema_title = annotation.schema_title
        else:
            schema_title = str(annotation.__name__)
    elif hasattr(annotation, '__forward_arg__'):
        schema_title = str(annotation.__forward_arg__)
    else:
        schema_title = str(annotation)
    type_struct = get_canonical_type_name_for_type(schema_title)
    return type_struct or schema_title


def is_typed_named_tuple_annotation(annotation: Any) -> bool:
    return hasattr(annotation, '_fields') and hasattr(annotation,
                                                      '__annotations__')
