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

import pydash
from marshmallow import EXCLUDE, INCLUDE, fields, post_dump, pre_load

from ..._schema import NestedField, StringTransformedEnum, UnionField
from ..._schema.component.component import ComponentSchema
from ..._schema.core.fields import ArmVersionedStr, CodeField, EnvironmentField, RegistryStr
from ..._schema.job.parameterized_spark import SparkEntryClassSchema, SparkEntryFileSchema
from ..._utils._arm_id_utils import parse_name_label
from ..._utils.utils import get_valid_dot_keys_with_wildcard
from ...constants._common import (
    LABELLED_RESOURCE_NAME,
    SOURCE_PATH_CONTEXT_KEY,
    AzureMLResourceType,
    DefaultOpenEncoding,
)
from ...constants._component import NodeType as PublicNodeType
from .._utils import yaml_safe_load_with_base_resolver
from .environment import InternalEnvironmentSchema
from .input_output import (
    InternalEnumParameterSchema,
    InternalInputPortSchema,
    InternalOutputPortSchema,
    InternalParameterSchema,
    InternalPrimitiveOutputSchema,
    InternalSparkParameterSchema,
)


class NodeType:
    COMMAND = "CommandComponent"
    DATA_TRANSFER = "DataTransferComponent"
    DISTRIBUTED = "DistributedComponent"
    HDI = "HDInsightComponent"
    SCOPE_V2 = "scope"
    HDI_V2 = "hdinsight"
    HEMERA_V2 = "hemera"
    STARLITE_V2 = "starlite"
    AE365EXEPOOL_V2 = "ae365exepool"
    AETHER_BRIDGE_V2 = "aetherbridge"
    PARALLEL = "ParallelComponent"
    SCOPE = "ScopeComponent"
    STARLITE = "StarliteComponent"
    SWEEP = "SweepComponent"
    PIPELINE = "PipelineComponent"
    HEMERA = "HemeraComponent"
    AE365EXEPOOL = "AE365ExePoolComponent"
    IPP = "IntellectualPropertyProtectedComponent"
    # internal spake component got a type value conflict with spark component
    # this enum is used to identify its create_function in factories
    SPARK = "DummySpark"
    AETHER_BRIDGE = "AetherBridgeComponent"

    @classmethod
    def all_values(cls):
        all_values = []
        for key, value in vars(cls).items():
            if not key.startswith("_") and isinstance(value, str):
                all_values.append(value)
        return all_values


class InternalComponentSchema(ComponentSchema):
    class Meta:
        unknown = INCLUDE

    # override name as 1p components allow . in name, which is not allowed in v2 components
    name = fields.Str()

    # override to allow empty properties
    tags = fields.Dict(keys=fields.Str())

    # override inputs & outputs to support 1P inputs & outputs, may need to do strict validation later
    # no need to check io type match since server will do that
    inputs = fields.Dict(
        keys=fields.Str(),
        values=UnionField(
            [
                NestedField(InternalParameterSchema),
                NestedField(InternalEnumParameterSchema),
                NestedField(InternalInputPortSchema),
            ]
        ),
    )
    # support primitive output for all internal components for now
    outputs = fields.Dict(
        keys=fields.Str(),
        values=UnionField(
            [
                NestedField(InternalPrimitiveOutputSchema, unknown=EXCLUDE),
                NestedField(InternalOutputPortSchema, unknown=EXCLUDE),
            ]
        ),
    )

    # type field is required for registration
    type = StringTransformedEnum(
        allowed_values=NodeType.all_values(),
        casing_transform=lambda x: parse_name_label(x)[0],
        pass_original=True,
    )

    # need to resolve as it can be a local field
    code = CodeField()

    environment = UnionField(
        [
            RegistryStr(azureml_type=AzureMLResourceType.ENVIRONMENT),
            ArmVersionedStr(azureml_type=AzureMLResourceType.ENVIRONMENT),
            NestedField(InternalEnvironmentSchema),
        ]
    )

    def get_skip_fields(self):
        return ["properties"]

    def _serialize(self, obj, *, many: bool = False):
        if many and obj is not None:
            return super(InternalComponentSchema, self)._serialize(obj, many=many)
        ret = super(InternalComponentSchema, self)._serialize(obj)
        for attr_name in obj.__dict__.keys():
            if (
                not attr_name.startswith("_")
                and attr_name not in self.get_skip_fields()
                and attr_name not in self.dump_fields
            ):
                ret[attr_name] = self.get_attribute(obj, attr_name, None)
        return ret

    # override param_override to ensure that param override happens after reloading the yaml
    @pre_load
    def add_param_overrides(self, data, **kwargs):
        source_path = self.context.pop(SOURCE_PATH_CONTEXT_KEY, None)
        if isinstance(data, dict) and source_path and os.path.isfile(source_path):

            def should_node_overwritten(_root, _parts):
                parts = _parts.copy()
                parts.pop()
                parts.append("type")
                _input_type = pydash.get(_root, parts, None)
                return isinstance(_input_type, str) and _input_type.lower() not in ["boolean"]

            # do override here
            with open(source_path, "r", encoding=DefaultOpenEncoding.READ) as f:
                origin_data = yaml_safe_load_with_base_resolver(f)
                for dot_key_wildcard, condition_func in [
                    ("version", None),
                    ("inputs.*.default", should_node_overwritten),
                    ("inputs.*.enum", should_node_overwritten),
                ]:
                    for dot_key in get_valid_dot_keys_with_wildcard(
                        origin_data, dot_key_wildcard, validate_func=condition_func
                    ):
                        pydash.set_(data, dot_key, pydash.get(origin_data, dot_key))
        return super().add_param_overrides(data, **kwargs)

    @post_dump(pass_original=True)
    def simplify_input_output_port(self, data, original, **kwargs):  # pylint:disable=unused-argument
        # remove None in input & output
        for io_ports in [data["inputs"], data["outputs"]]:
            for port_name, port_definition in io_ports.items():
                io_ports[port_name] = dict(filter(lambda item: item[1] is not None, port_definition.items()))

        # hack, to match current serialization match expectation
        for port_name, port_definition in data["inputs"].items():
            if "mode" in port_definition:
                del port_definition["mode"]

        return data

    @post_dump(pass_original=True)
    def add_back_type_label(self, data, original, **kwargs):  # pylint:disable=unused-argument
        type_label = original._type_label  # pylint:disable=protected-access
        if type_label:
            data["type"] = LABELLED_RESOURCE_NAME.format(data["type"], type_label)
        return data


class InternalSparkComponentSchema(InternalComponentSchema):
    # type field is required for registration
    type = StringTransformedEnum(
        allowed_values=PublicNodeType.SPARK,
        casing_transform=lambda x: parse_name_label(x)[0].lower(),
        pass_original=True,
    )

    # override inputs:
    # https://componentsdk.azurewebsites.net/components/spark_component.html#differences-with-other-component-types
    inputs = fields.Dict(
        keys=fields.Str(),
        values=UnionField(
            [
                NestedField(InternalSparkParameterSchema),
                NestedField(InternalInputPortSchema),
            ]
        ),
    )

    environment = EnvironmentField(
        extra_fields=[NestedField(InternalEnvironmentSchema)],
        allow_none=True,
    )

    jars = UnionField(
        [
            fields.List(fields.Str()),
            fields.Str(),
        ],
    )
    py_files = UnionField(
        [
            fields.List(fields.Str()),
            fields.Str(),
        ],
        data_key="pyFiles",
        attribute="py_files",
    )

    entry = UnionField(
        [NestedField(SparkEntryFileSchema), NestedField(SparkEntryClassSchema)],
        required=True,
        metadata={"description": "Entry."},
    )

    files = fields.List(fields.Str(required=True))
    archives = fields.List(fields.Str(required=True))
    conf = fields.Dict(keys=fields.Str(), values=fields.Raw())
    args = fields.Str(metadata={"description": "Command Line arguments."})
