# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
"""The preprocessing autorest plugin."""

import copy
from typing import Callable, Any, Optional

from ..utils import to_snake_case, extract_original_name
from .helpers import (
    add_redefined_builtin_info,
    pad_builtin_namespaces,
    pad_special_chars,
)
from .python_mappings import TSP_RESERVED_WORDS, RESERVED_WORDS, PadType

from .. import YamlUpdatePlugin
from ..utils import (
    parse_args,
    get_body_type_for_description,
    JSON_REGEXP,
    KNOWN_TYPES,
)


def update_overload_section(
    overload: dict[str, Any],
    yaml_data: dict[str, Any],
    section: str,
):
    try:
        for overload_s, original_s in zip(overload[section], yaml_data[section]):
            if overload_s.get("type"):
                overload_s["type"] = original_s["type"]
            if overload_s.get("headers"):
                for overload_h, original_h in zip(overload_s["headers"], original_s["headers"]):
                    if overload_h.get("type"):
                        overload_h["type"] = original_h["type"]
    except KeyError as exc:
        raise ValueError(overload["name"]) from exc


def add_overload(yaml_data: dict[str, Any], body_type: dict[str, Any], for_flatten_params=False):
    overload = copy.deepcopy(yaml_data)
    overload["isOverload"] = True
    overload["bodyParameter"]["type"] = body_type
    overload["bodyParameter"]["defaultToUnsetSentinel"] = False
    overload["overloads"] = []
    if yaml_data.get("initialOperation"):
        overload["initialOperation"] = yaml_data["initialOperation"]

    if for_flatten_params:
        overload["bodyParameter"]["flattened"] = True
    else:
        overload["parameters"] = [p for p in overload["parameters"] if not p.get("inFlattenedBody")]
    # for yaml sync, we need to make sure all of the responses, parameters, and exceptions' types have the same yaml id
    for overload_p, original_p in zip(overload["parameters"], yaml_data["parameters"]):
        overload_p["type"] = original_p["type"]
    update_overload_section(overload, yaml_data, "responses")
    update_overload_section(overload, yaml_data, "exceptions")

    # update content type to be an overloads content type
    content_type_param = next(p for p in overload["parameters"] if p["wireName"].lower() == "content-type")
    content_type_param["inOverload"] = True
    content_type_param["inDocstring"] = True
    body_type_description = get_body_type_for_description(overload["bodyParameter"])
    content_type_param["description"] = (
        f"Body Parameter content-type. Content type parameter for {body_type_description} body."
    )
    content_types = yaml_data["bodyParameter"]["contentTypes"]
    if body_type["type"] == "binary" and len(content_types) > 1:
        content_types = "'" + "', '".join(content_types) + "'"
        content_type_param["description"] += f" Known values are: {content_types}."
    overload["bodyParameter"]["inOverload"] = True
    for parameter in overload["parameters"]:
        parameter["inOverload"] = True
        parameter["defaultToUnsetSentinel"] = False
    return overload


def add_overloads_for_body_param(yaml_data: dict[str, Any]) -> None:
    """If we added a body parameter type, add overloads for that type"""
    body_parameter = yaml_data["bodyParameter"]
    if not (
        body_parameter["type"]["type"] == "combined"
        and len(yaml_data["bodyParameter"]["type"]["types"]) > len(yaml_data["overloads"])
    ):
        return
    for body_type in body_parameter["type"]["types"]:
        if any(o for o in yaml_data["overloads"] if id(o["bodyParameter"]["type"]) == id(body_type)):
            continue
        if body_type.get("type") == "model" and body_type.get("base") == "json":
            yaml_data["overloads"].append(add_overload(yaml_data, body_type, for_flatten_params=True))
        yaml_data["overloads"].append(add_overload(yaml_data, body_type))
    content_type_param = next(p for p in yaml_data["parameters"] if p["wireName"].lower() == "content-type")
    content_type_param["inOverload"] = False
    content_type_param["inOverridden"] = True
    content_type_param["inDocstring"] = True
    content_type_param["clientDefaultValue"] = (
        None  # make it none bc it will be overridden, we depend on default of overloads
    )
    content_type_param["optional"] = True


def update_description(description: Optional[str], default_description: str = "") -> str:
    if not description:
        description = default_description
    description.rstrip(" ")
    if description and description[-1] != ".":
        description += "."
    return description


def update_operation_group_class_name(prefix: str, class_name: str) -> str:
    if class_name == "":
        return prefix + "OperationsMixin"
    if class_name == "Operations":
        return "Operations"
    return class_name + "Operations"


def update_paging_response(yaml_data: dict[str, Any]) -> None:
    yaml_data["discriminator"] = "paging"


HEADERS_HIDE_IN_METHOD = (
    "repeatability-request-id",
    "repeatability-first-sent",
    "x-ms-client-request-id",
    "client-request-id",
    "return-client-request-id",
)
HEADERS_CONVERT_IN_METHOD = {
    "if-match": {
        "clientName": "etag",
        "wireName": "etag",
        "description": "check if resource is changed. Set None to skip checking etag.",
    },
    "if-none-match": {
        "clientName": "match_condition",
        "wireName": "match-condition",
        "description": "The match condition to use upon the etag.",
        "type": {
            "type": "sdkcore",
            "name": "MatchConditions",
        },
    },
}
CLOUD_SETTING = {
    "optional": True,
    "description": "The cloud setting for which to get the ARM endpoint.",
    "clientName": "cloud_setting",
    "implementation": "Client",
    "location": "keyword",
    "type": {
        "type": "sdkcore",
        "name": "AzureClouds",
        "isTypingOnly": True,
    },
}


def get_wire_name_lower(parameter: dict[str, Any]) -> str:
    return (parameter.get("wireName") or "").lower()


def headers_convert(yaml_data: dict[str, Any], replace_data: Any) -> None:
    if isinstance(replace_data, dict):
        for k, v in replace_data.items():
            yaml_data[k] = v


def has_json_content_type(yaml_data: dict[str, Any]) -> bool:
    return any(ct for ct in yaml_data.get("contentTypes", []) if JSON_REGEXP.match(ct))


def has_multi_part_content_type(yaml_data: dict[str, Any]) -> bool:
    return any(ct for ct in yaml_data.get("contentTypes", []) if ct == "multipart/form-data")


class PreProcessPlugin(YamlUpdatePlugin):
    """Add Python naming information."""

    @property
    def azure_arm(self) -> bool:
        return self.options.get("azure-arm", False)

    @property
    def version_tolerant(self) -> bool:
        return self.options.get("version-tolerant", True)

    @property
    def models_mode(self) -> Optional[str]:
        return self.options.get("models-mode", "dpg" if self.is_tsp else None)

    @property
    def is_tsp(self) -> bool:
        return self.options.get("tsp_file", False)

    def add_body_param_type(
        self,
        code_model: dict[str, Any],
        body_parameter: dict[str, Any],
    ):
        # only add overload for special content type
        if (  # pylint: disable=too-many-boolean-expressions
            body_parameter
            and body_parameter["type"]["type"] in ("model", "dict", "list")
            and (has_json_content_type(body_parameter) or (self.is_tsp and has_multi_part_content_type(body_parameter)))
            and not body_parameter["type"].get("xmlMetadata")
            and not any(t for t in ["flattened", "groupedBy"] if body_parameter.get(t))
        ):
            origin_type = body_parameter["type"]["type"]
            model_type = (
                body_parameter["type"] if origin_type == "model" else body_parameter["type"].get("elementType", {})
            )
            is_dpg_model = model_type.get("base") == "dpg"
            body_parameter["type"] = {
                "type": "combined",
                "types": [body_parameter["type"]],
            }
            # don't add binary overload for multipart content type
            if not (self.is_tsp and has_multi_part_content_type(body_parameter)):
                body_parameter["type"]["types"].append(KNOWN_TYPES["binary"])

            if self.options["models-mode"] == "dpg" and is_dpg_model:
                if origin_type == "model":
                    body_parameter["type"]["types"].insert(1, KNOWN_TYPES["any-object"])
                else:
                    # dict or list
                    # copy the original dict / list type
                    any_obj_list_or_dict = copy.deepcopy(body_parameter["type"]["types"][0])
                    any_obj_list_or_dict["elementType"] = KNOWN_TYPES["any-object"]
                    body_parameter["type"]["types"].insert(1, any_obj_list_or_dict)
            code_model["types"].append(body_parameter["type"])

    def pad_reserved_words(self, name: str, pad_type: PadType, yaml_type: dict[str, Any]) -> str:
        # we want to pad hidden variables as well
        if not name:
            # we'll pass in empty operation groups sometime etc.
            return name

        if self.is_tsp:
            reserved_words = {k: (v + TSP_RESERVED_WORDS.get(k, [])) for k, v in RESERVED_WORDS.items()}
        else:
            reserved_words = RESERVED_WORDS
        name = pad_special_chars(name)
        name_prefix = "_" if name[0] == "_" else ""
        name = name[1:] if name[0] == "_" else name
        if name.lower() in reserved_words[pad_type]:
            if self.is_tsp and name.lower() in TSP_RESERVED_WORDS.get(pad_type, []):
                # to maintain backcompat for cases where we pad in tsp but not in autorest,
                # if we have a tsp reserved word, we also want to keep track of the original name for backcompat
                yaml_type["originalTspName"] = name_prefix + name
            return name_prefix + name + pad_type
        return name_prefix + name

    def update_types(self, yaml_data: list[dict[str, Any]]) -> None:
        for type in yaml_data:
            for property in type.get("properties", []):
                property["description"] = update_description(property.get("description", ""))
                property["clientName"] = self.pad_reserved_words(
                    property["clientName"].lower(), PadType.PROPERTY, property
                )
                add_redefined_builtin_info(property["clientName"], property)
            if type.get("name"):
                pad_type = PadType.MODEL if type["type"] == "model" else PadType.ENUM_CLASS
                if type["type"] != "enumvalue":
                    name = self.pad_reserved_words(type["name"], pad_type, type)
                    type["name"] = name[0].upper() + name[1:]
                type["description"] = update_description(type.get("description", ""), type["name"])
                type["snakeCaseName"] = to_snake_case(type["name"])
            if type.get("values"):
                # we're enums - enum values are UPPER_CASE so no padding needed for reserved words
                for value in type["values"]:
                    upper_name = value["name"].upper()
                    if upper_name[0] in "0123456789":
                        upper_name = "ENUM_" + upper_name
                        value["name"] = upper_name

        # add type for reference
        for v in HEADERS_CONVERT_IN_METHOD.values():
            if isinstance(v, dict) and "type" in v:
                yaml_data.append(v["type"])
        yaml_data.append(CLOUD_SETTING["type"])  # type: ignore

    def update_client(self, yaml_data: dict[str, Any]) -> None:
        yaml_data["description"] = update_description(yaml_data["description"], default_description=yaml_data["name"])
        yaml_data["legacyFilename"] = to_snake_case(yaml_data["name"].replace(" ", "_"))
        parameters = yaml_data["parameters"]
        for parameter in parameters:
            self.update_parameter(parameter)
            if parameter["clientName"] == "credential":
                policy = parameter["type"].get("policy")
                if policy and policy["type"] == "BearerTokenCredentialPolicy" and self.azure_arm:
                    policy["type"] = "ARMChallengeAuthenticationPolicy"
                    policy["credentialScopes"] = ["https://management.azure.com/.default"]
        if (
            (not self.version_tolerant or self.azure_arm)
            and parameters
            and parameters[-1]["clientName"] == "credential"
        ):
            # we need to move credential to the front in mgmt mode for backcompat reasons
            yaml_data["parameters"] = [parameters[-1]] + parameters[:-1]
        prop_name = yaml_data["name"]
        if prop_name.endswith("Client"):
            prop_name = prop_name[: len(prop_name) - len("Client")]
        yaml_data["builderPadName"] = to_snake_case(prop_name)
        for og in yaml_data.get("operationGroups", []):
            for o in og["operations"]:
                property_if_match = None
                property_if_none_match = None
                for p in o["parameters"]:
                    wire_name_lower = get_wire_name_lower(p)
                    if p["location"] == "header" and wire_name_lower == "client-request-id":
                        yaml_data["requestIdHeaderName"] = wire_name_lower
                    if self.version_tolerant and p["location"] == "header":
                        if wire_name_lower == "if-match":
                            property_if_match = p
                        elif wire_name_lower == "if-none-match":
                            property_if_none_match = p
                # pylint: disable=line-too-long
                # some service(e.g. https://github.com/Azure/azure-rest-api-specs/blob/main/specification/cosmos-db/data-plane/Microsoft.Tables/preview/2019-02-02/table.json)
                # only has one, so we need to add "if-none-match" or "if-match" if it's missing
                if not property_if_match and property_if_none_match:
                    property_if_match = property_if_none_match.copy()
                    property_if_match["wireName"] = "if-match"
                if not property_if_none_match and property_if_match:
                    property_if_none_match = property_if_match.copy()
                    property_if_none_match["wireName"] = "if-none-match"

                if property_if_match and property_if_none_match:
                    # arrange if-match and if-none-match to the end of parameters
                    o["parameters"] = [
                        item
                        for item in o["parameters"]
                        if get_wire_name_lower(item) not in ("if-match", "if-none-match")
                    ] + [property_if_match, property_if_none_match]

                    o["hasEtag"] = True
                    yaml_data["hasEtag"] = True

        # add client signature cloud_setting for arm
        if self.azure_arm and yaml_data["parameters"]:
            yaml_data["parameters"].append(CLOUD_SETTING)

    def get_operation_updater(self, yaml_data: dict[str, Any]) -> Callable[[dict[str, Any], dict[str, Any]], None]:
        if yaml_data["discriminator"] == "lropaging":
            return self.update_lro_paging_operation
        if yaml_data["discriminator"] == "lro":
            return self.update_lro_operation
        if yaml_data["discriminator"] == "paging":
            return self.update_paging_operation
        return self.update_operation

    def update_parameter(self, yaml_data: dict[str, Any]) -> None:
        yaml_data["description"] = update_description(yaml_data.get("description", ""))
        if not (yaml_data["location"] == "header" and yaml_data["clientName"] in ("content_type", "accept")):
            yaml_data["clientName"] = self.pad_reserved_words(
                yaml_data["clientName"].lower(), PadType.PARAMETER, yaml_data
            )
        if yaml_data.get("propertyToParameterName"):
            # need to create a new one with padded values (but NOT keys, since keys are wire names)
            yaml_data["propertyToParameterName"] = {
                prop: self.pad_reserved_words(param_name, PadType.PARAMETER, yaml_data).lower()
                for prop, param_name in yaml_data["propertyToParameterName"].items()
            }
        wire_name_lower = (yaml_data.get("wireName") or "").lower()
        if yaml_data["location"] == "header" and (
            wire_name_lower in HEADERS_HIDE_IN_METHOD or yaml_data.get("clientDefaultValue") == "multipart/form-data"
        ):
            yaml_data["hideInMethod"] = True
        if self.version_tolerant and yaml_data["location"] == "header" and wire_name_lower in HEADERS_CONVERT_IN_METHOD:
            headers_convert(yaml_data, HEADERS_CONVERT_IN_METHOD[wire_name_lower])
        if wire_name_lower in ["$host", "content-type", "accept"] and yaml_data["type"]["type"] == "constant":
            yaml_data["clientDefaultValue"] = yaml_data["type"]["value"]

    def update_operation(
        self,
        code_model: dict[str, Any],
        yaml_data: dict[str, Any],
        *,
        is_overload: bool = False,
    ) -> None:
        yaml_data["groupName"] = self.pad_reserved_words(yaml_data["groupName"], PadType.OPERATION_GROUP, yaml_data)
        yaml_data["groupName"] = to_snake_case(yaml_data["groupName"])
        yaml_data["name"] = yaml_data["name"].lower()
        if yaml_data.get("isLroInitialOperation") is True:
            yaml_data["name"] = (
                "_"
                + self.pad_reserved_words(extract_original_name(yaml_data["name"]), PadType.METHOD, yaml_data)
                + "_initial"
            )
        else:
            yaml_data["name"] = self.pad_reserved_words(yaml_data["name"], PadType.METHOD, yaml_data)
        yaml_data["description"] = update_description(yaml_data["description"], yaml_data["name"])
        yaml_data["summary"] = update_description(yaml_data.get("summary", ""))
        body_parameter = yaml_data.get("bodyParameter")
        for parameter in yaml_data["parameters"]:
            self.update_parameter(parameter)
        if yaml_data.get("bodyParameter"):
            self.update_parameter(yaml_data["bodyParameter"])
            for entry in yaml_data["bodyParameter"].get("entries", []):
                self.update_parameter(entry)
        for overload in yaml_data.get("overloads", []):
            self.update_operation(code_model, overload, is_overload=True)
        for response in yaml_data.get("responses", []):
            response["discriminator"] = "operation"
        if body_parameter and not is_overload:
            # if we have a JSON body, we add a binary overload
            self.add_body_param_type(code_model, body_parameter)
            add_overloads_for_body_param(yaml_data)

    def _update_lro_operation_helper(self, yaml_data: dict[str, Any]) -> None:
        for response in yaml_data.get("responses", []):
            response["discriminator"] = "lro"
            response["pollerSync"] = response.get("pollerSync") or "azure.core.polling.LROPoller"
            response["pollerAsync"] = response.get("pollerAsync") or "azure.core.polling.AsyncLROPoller"
            if not response.get("pollingMethodSync"):
                response["pollingMethodSync"] = (
                    "azure.mgmt.core.polling.arm_polling.ARMPolling"
                    if self.azure_arm
                    else "azure.core.polling.base_polling.LROBasePolling"
                )
            if not response.get("pollingMethodAsync"):
                response["pollingMethodAsync"] = (
                    "azure.mgmt.core.polling.async_arm_polling.AsyncARMPolling"
                    if self.azure_arm
                    else "azure.core.polling.async_base_polling.AsyncLROBasePolling"
                )

    def update_lro_paging_operation(
        self,
        code_model: dict[str, Any],
        yaml_data: dict[str, Any],
        is_overload: bool = False,
        item_type: Optional[dict[str, Any]] = None,
    ) -> None:
        self.update_lro_operation(code_model, yaml_data, is_overload=is_overload)
        self.update_paging_operation(code_model, yaml_data, is_overload=is_overload, item_type=item_type)
        yaml_data["discriminator"] = "lropaging"
        for response in yaml_data.get("responses", []):
            response["discriminator"] = "lropaging"
        for overload in yaml_data.get("overloads", []):
            self.update_lro_paging_operation(
                code_model,
                overload,
                is_overload=True,
                item_type=yaml_data["responses"][0]["itemType"],
            )

    def update_lro_operation(
        self,
        code_model: dict[str, Any],
        yaml_data: dict[str, Any],
        is_overload: bool = False,
    ) -> None:
        def convert_initial_operation_response_type(data: dict[str, Any]) -> None:
            for response in data.get("responses", []):
                response["type"] = KNOWN_TYPES["binary"]

        self.update_operation(code_model, yaml_data, is_overload=is_overload)
        self.update_operation(code_model, yaml_data["initialOperation"], is_overload=is_overload)
        convert_initial_operation_response_type(yaml_data["initialOperation"])
        self._update_lro_operation_helper(yaml_data)
        for overload in yaml_data.get("overloads", []):
            self._update_lro_operation_helper(overload)
            self.update_operation(code_model, overload["initialOperation"], is_overload=True)
            convert_initial_operation_response_type(overload["initialOperation"])

    def update_paging_operation(
        self,
        code_model: dict[str, Any],
        yaml_data: dict[str, Any],
        is_overload: bool = False,
        item_type: Optional[dict[str, Any]] = None,
    ) -> None:
        self.update_operation(code_model, yaml_data, is_overload=is_overload)
        item_type = item_type or yaml_data["itemType"]["elementType"]
        if yaml_data.get("nextOperation"):
            yaml_data["nextOperation"]["groupName"] = self.pad_reserved_words(
                yaml_data["nextOperation"]["groupName"], PadType.OPERATION_GROUP, yaml_data["nextOperation"]
            )
            yaml_data["nextOperation"]["groupName"] = to_snake_case(yaml_data["nextOperation"]["groupName"])
            for response in yaml_data["nextOperation"].get("responses", []):
                update_paging_response(response)
                response["itemType"] = item_type
        for response in yaml_data.get("responses", []):
            update_paging_response(response)
            response["itemType"] = item_type
        for overload in yaml_data.get("overloads", []):
            self.update_paging_operation(code_model, overload, is_overload=True, item_type=item_type)

    def update_operation_groups(self, code_model: dict[str, Any], client: dict[str, Any]) -> None:
        operation_groups_yaml_data = client.get("operationGroups", [])
        for operation_group in operation_groups_yaml_data:
            operation_group["identifyName"] = self.pad_reserved_words(
                operation_group.get("name", operation_group["propertyName"]),
                PadType.OPERATION_GROUP,
                operation_group,
            )
            operation_group["identifyName"] = to_snake_case(operation_group["identifyName"])
            operation_group["propertyName"] = self.pad_reserved_words(
                operation_group["propertyName"], PadType.OPERATION_GROUP, operation_group
            )
            operation_group["propertyName"] = to_snake_case(operation_group["propertyName"])
            operation_group["className"] = update_operation_group_class_name(
                client["name"], operation_group["className"]
            )
            for operation in operation_group["operations"]:
                self.get_operation_updater(operation)(code_model, operation)

            if operation_group.get("operationGroups"):
                self.update_operation_groups(code_model, operation_group)

    def update_yaml(self, yaml_data: dict[str, Any]) -> None:
        """Convert in place the YAML str."""
        self.update_types(yaml_data["types"])
        yaml_data["types"] += KNOWN_TYPES.values()
        for client in yaml_data["clients"]:
            self.update_client(client)
            self.update_operation_groups(yaml_data, client)
        if yaml_data.get("namespace"):
            yaml_data["namespace"] = pad_builtin_namespaces(yaml_data["namespace"])


if __name__ == "__main__":
    # TSP pipeline will call this
    args, unknown_args = parse_args()
    PreProcessPlugin(output_folder=args.output_folder, tsp_file=args.tsp_file, **unknown_args).process()
