# ---------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# ---------------------------------------------------------

# pylint: disable=protected-access
# enable protected access for protected helper functions

import copy
from collections import OrderedDict
from enum import Enum as PyEnum
from enum import EnumMeta
from inspect import Parameter, signature

from azure.ai.ml.constants._component import IOConstants
from azure.ai.ml.exceptions  import UserErrorException


def is_parameter_group(obj):
    """Return True if obj is a parameter group or an instance of a parameter group class."""
    return hasattr(obj, IOConstants.GROUP_ATTR_NAME)


def _get_annotation_by_value(val):
    # TODO: we'd better remove this potential recursive import
    from .enum_input import EnumInput
    from .input import Input

    def _is_dataset(data):
        from azure.ai.ml.entities._job.job_io_mixin import JobIOMixin

        DATASET_TYPES = JobIOMixin
        return isinstance(data, DATASET_TYPES)

    if _is_dataset(val):
        annotation = Input
    elif val is Parameter.empty or val is None:
        # If no default value or default is None, create val as the basic parameter type,
        # it could be replaced using component parameter definition.
        annotation = Input._get_default_unknown_input()
    elif isinstance(val, PyEnum):
        # Handle enum values
        annotation = EnumInput(enum=val.__class__)
    else:
        annotation = _get_annotation_cls_by_type(type(val), raise_error=False)
        if not annotation:
            # Fall back to default
            annotation = Input._get_default_unknown_input()
    return annotation


def _get_annotation_cls_by_type(t: type, raise_error=False, optional=None):
    # TODO: we'd better remove this potential recursive import
    from .input import Input

    cls = Input._get_input_by_type(t, optional=optional)
    if cls is None and raise_error:
        raise UserErrorException(f"Can't convert type {t} to azure.ai.ml.Input")
    return cls


# pylint: disable=too-many-statements
def _get_param_with_standard_annotation(cls_or_func, is_func=False, skip_params=None):
    """Standardize function parameters or class fields with dsl.types
    annotation."""
    # TODO: we'd better remove this potential recursive import
    from .input import Input
    from .output import Output

    def _is_dsl_type_cls(t: type):
        if type(t) is not type:  # pylint: disable=unidiomatic-typecheck
            return False
        return issubclass(t, (Input, Output))

    def _is_dsl_types(o: object):
        return _is_dsl_type_cls(type(o))

    def _get_fields(annotations):
        """Return field names to annotations mapping in class."""
        annotation_fields = OrderedDict()
        for name, annotation in annotations.items():
            # Skip return type
            if name == "return":
                continue
            # Handle EnumMeta annotation
            if isinstance(annotation, EnumMeta):
                from .enum_input import EnumInput

                annotation = EnumInput(type="string", enum=annotation)
            # Handle Group annotation
            if is_parameter_group(annotation):
                annotation = copy.deepcopy(getattr(annotation, IOConstants.GROUP_ATTR_NAME))
            # Try creating annotation by type when got like 'param: int'
            if not _is_dsl_type_cls(annotation) and not _is_dsl_types(annotation):
                origin_annotation = annotation
                annotation = _get_annotation_cls_by_type(annotation, raise_error=False)
                if not annotation:
                    msg = f"Unsupported annotation type {origin_annotation!r} for parameter {name!r}."
                    raise UserErrorException(msg)
            annotation_fields[name] = annotation
        return annotation_fields

    def _merge_field_keys(annotation_fields, defaults_dict):
        """Merge field keys from annotations and cls dict to get all fields in
        class."""
        anno_keys = list(annotation_fields.keys())
        dict_keys = defaults_dict.keys()
        if not dict_keys:
            return anno_keys
        return [*anno_keys, *[key for key in dict_keys if key not in anno_keys]]

    def _update_annotation_with_default(anno, name, default):
        """Create annotation if is type class and update the default."""
        # Create instance if is type class
        complete_annotation = anno
        if _is_dsl_type_cls(anno):
            complete_annotation = anno()
        complete_annotation.name = name
        if default is Input._EMPTY:
            return complete_annotation
        if isinstance(complete_annotation, Input):
            # Non-parameter Input has no default attribute
            if complete_annotation._is_primitive_type and complete_annotation.default is not None:
                # logger.warning(
                #     f"Warning: Default value of f{complete_annotation.name!r} is set twice: "
                #     f"{complete_annotation.default!r} and {default!r}, will use {default!r}"
                # )
                pass
            complete_annotation._update_default(default)
        return complete_annotation

    def _update_fields_with_default(annotation_fields, defaults_dict):
        """Use public values in class dict to update annotations."""
        all_fields = OrderedDict()
        all_filed_keys = _merge_field_keys(annotation_fields, defaults_dict)
        for name in all_filed_keys:
            # Get or create annotation
            annotation = (
                annotation_fields[name]
                if name in annotation_fields
                else _get_annotation_by_value(defaults_dict.get(name, Input._EMPTY))
            )
            # Create annotation if is class type and update default
            annotation = _update_annotation_with_default(annotation, name, defaults_dict.get(name, Input._EMPTY))
            all_fields[name] = annotation
        return all_fields

    def _get_inherited_fields():
        """Get all fields inherit from bases parameter group class."""
        _fields = OrderedDict({})
        if is_func:
            return _fields
        # In reversed order so that more derived classes
        # override earlier field definitions in base classes.
        for base in cls_or_func.__mro__[-1:0:-1]:
            if is_parameter_group(base):
                # merge and reorder fields from current base with previous
                _fields = _merge_and_reorder(_fields, copy.deepcopy(getattr(base, IOConstants.GROUP_ATTR_NAME).values))
        return _fields

    def _merge_and_reorder(inherited_fields, cls_fields):
        """Merge inherited fields with cls fields. The order inside each part
        will not be changed. Order will be.

        {inherited_no_default_fields} + {cls_no_default_fields} + {inherited_default_fields} + {cls_default_fields}.
        Note: If cls overwrite an inherited no default field with default, it will be put in the
        cls_default_fields part and deleted from inherited_no_default_fields:
        e.g.
        @dsl.parameter_group
        class SubGroup:
            int_param0: Integer
            int_param1: int
        @dsl.parameter_group
        class Group(SubGroup):
            int_param3: Integer
            int_param1: int = 1
        The init function of Group will be 'def __init__(self, *, int_param0, int_param3, int_param1=1)'.
        """

        def _split(_fields):
            """Split fields to two parts from the first default field."""
            _no_defaults_fields, _defaults_fields = {}, {}
            seen_default = False
            for key, val in _fields.items():
                if val.get("default", None) or seen_default:
                    seen_default = True
                    _defaults_fields[key] = val
                else:
                    _no_defaults_fields[key] = val
            return _no_defaults_fields, _defaults_fields

        inherited_no_default, inherited_default = _split(inherited_fields)
        cls_no_default, cls_default = _split(cls_fields)
        # Cross comparison and delete from inherited_fields if same key appeared in cls_fields
        # pylint: disable=consider-iterating-dictionary
        for key in cls_default.keys():
            if key in inherited_no_default.keys():
                del inherited_no_default[key]
        for key in cls_no_default.keys():
            if key in inherited_default.keys():
                del inherited_default[key]
        return OrderedDict(
            {
                **inherited_no_default,
                **cls_no_default,
                **inherited_default,
                **cls_default,
            }
        )

    skip_params = skip_params or []
    inherited_fields = _get_inherited_fields()
    # From annotations get field with type
    annotations = getattr(cls_or_func, "__annotations__", {})
    annotations = _update_io_from_mldesigner(annotations)
    annotation_fields = _get_fields(annotations)
    # Update fields use class field with defaults from class dict or signature(func).paramters
    if not is_func:
        # Only consider public fields in class dict
        defaults_dict = {
            key: val for key, val in cls_or_func.__dict__.items()
            if not key.startswith("_") and key not in skip_params
        }
    else:
        # Infer parameter type from value if is function
        defaults_dict = {
            key: val.default for key, val in signature(cls_or_func).parameters.items()
            if key not in skip_params
        }
    fields = _update_fields_with_default(annotation_fields, defaults_dict)
    all_fields = _merge_and_reorder(inherited_fields, fields)
    return all_fields


def _update_io_from_mldesigner(annotations: dict) -> dict:
    """This function will translate IOBase from mldesigner package to azure.ml.entities.Input/Output.

    This function depend on `mldesigner._input_output._IOBase._to_io_entity_args_dict` to translate Input/Output
    instance annotations to IO entities.
    This function depend on class names of `mldesigner._input_output` to translate Input/Output class annotations
    to IO entities.
    """
    from azure.ai.ml import Input, Output
    mldesigner_pkg = "mldesigner"

    def _is_input_or_output_type(io: type, type_str: str):
        """Return true if type name contains type_str"""
        if isinstance(io, type) and io.__module__.startswith(mldesigner_pkg):
            if type_str in io.__name__:
                return True
        return False

    result = {}
    for key, io in annotations.items():
        if isinstance(io, type):
            if _is_input_or_output_type(io, "Input"):
                # mldesigner.Input -> entities.Input
                io = Input
            elif _is_input_or_output_type(io, "Output"):
                # mldesigner.Output -> entities.Output
                io = Output
        elif hasattr(io, "_to_io_entity_args_dict"):
            try:
                if _is_input_or_output_type(type(io), "Input"):
                    # mldesigner.Input() -> entities.Input()
                    io = Input(**io._to_io_entity_args_dict())
                elif _is_input_or_output_type(type(io), "Output"):
                    # mldesigner.Output() -> entities.Output()
                    io = Output(**io._to_io_entity_args_dict())
            except BaseException as e:
                raise UserErrorException(f"Failed to parse {io} to azure-ai-ml Input/Output: {str(e)}") from e
        result[key] = io
    return result


def _remove_empty_values(data, ignore_keys=None):
    if not isinstance(data, dict):
        return data
    ignore_keys = ignore_keys or {}
    return {
        k: v if k in ignore_keys else _remove_empty_values(v)
        for k, v in data.items()
        if v is not None or k in ignore_keys
    }
