# --------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
#
# Code generated by aaz-dev-tools
# --------------------------------------------------------------------------------------------

# pylint: skip-file
# flake8: noqa

from azure.cli.core.aaz import *


@register_command(
    "consumption budget update",
    is_preview=True,
)
class Update(AAZCommand):
    """Update operation to create or update a budget. Update operation requires latest eTag to be set in the request mandatorily. You may obtain the latest eTag by performing a get operation. Create operation does not require eTag.
    """

    _aaz_info = {
        "version": "2023-05-01",
        "resources": [
            ["mgmt-plane", "/subscriptions/{}/providers/microsoft.consumption/budgets/{}", "2023-05-01"],
        ]
    }

    AZ_SUPPORT_GENERIC_UPDATE = True

    def _handler(self, command_args):
        super()._handler(command_args)
        self._execute_operations()
        return self._output()

    _args_schema = None

    @classmethod
    def _build_arguments_schema(cls, *args, **kwargs):
        if cls._args_schema is not None:
            return cls._args_schema
        cls._args_schema = super()._build_arguments_schema(*args, **kwargs)

        # define Arg Group ""

        _args_schema = cls._args_schema
        _args_schema.budget_name = AAZStrArg(
            options=["-n", "--name", "--budget-name"],
            help="Budget Name.",
            required=True,
            id_part="name",
        )

        # define Arg Group "Parameters"

        _args_schema = cls._args_schema
        _args_schema.e_tag = AAZStrArg(
            options=["--e-tag"],
            arg_group="Parameters",
            help="eTag of the resource. To handle concurrent update scenario, this field will be used to determine whether the user is updating the latest version or not.",
            nullable=True,
        )

        # define Arg Group "Properties"

        _args_schema = cls._args_schema
        _args_schema.amount = AAZFloatArg(
            options=["--amount"],
            arg_group="Properties",
            help="The total amount of cost to track with the budget",
        )
        _args_schema.category = AAZStrArg(
            options=["--category"],
            arg_group="Properties",
            help="The category of the budget, whether the budget tracks cost or usage.",
            enum={"Cost": "Cost", "Usage": "Usage"},
        )
        _args_schema.filters = AAZObjectArg(
            options=["--filters"],
            arg_group="Properties",
            help="May be used to filter budgets by resource group, resource, or meter.",
            nullable=True,
        )
        _args_schema.notifications = AAZDictArg(
            options=["--notifications"],
            arg_group="Properties",
            help="Dictionary of notifications associated with the budget. Budget can have up to five notifications.",
            nullable=True,
        )
        _args_schema.time_grain = AAZStrArg(
            options=["--time-grain"],
            arg_group="Properties",
            help="The time covered by a budget. Tracking of the amount will be reset based on the time grain.",
            enum={"Annually": "Annually", "Monthly": "Monthly", "Quarterly": "Quarterly"},
        )
        _args_schema.time_period = AAZObjectArg(
            options=["--time-period"],
            arg_group="Properties",
            help="Has start and end date of the budget. The start date must be first of the month and should be less than the end date. Budget start date must be on or after June 1, 2017. Future start date should not be more than three months. Past start date should  be selected within the timegrain period. There are no restrictions on the end date.",
        )

        filters = cls._args_schema.filters
        filters.meters = AAZListArg(
            options=["meters"],
            help="The list of filters on meters, mandatory for budgets of usage category. ",
            nullable=True,
        )
        filters.resource_groups = AAZListArg(
            options=["resource-groups"],
            help="The list of filters on resource groups, allowed at subscription level only.",
            nullable=True,
        )
        filters.resources = AAZListArg(
            options=["resources"],
            help="The list of filters on resources.",
            nullable=True,
        )

        meters = cls._args_schema.filters.meters
        meters.Element = AAZStrArg(
            nullable=True,
        )

        resource_groups = cls._args_schema.filters.resource_groups
        resource_groups.Element = AAZStrArg(
            nullable=True,
        )

        resources = cls._args_schema.filters.resources
        resources.Element = AAZStrArg(
            nullable=True,
        )

        notifications = cls._args_schema.notifications
        notifications.Element = AAZObjectArg(
            nullable=True,
        )

        _element = cls._args_schema.notifications.Element
        _element.contact_emails = AAZListArg(
            options=["contact-emails"],
            help="Email addresses to send the budget notification to when the threshold is exceeded.",
        )
        _element.contact_groups = AAZListArg(
            options=["contact-groups"],
            help="Action groups to send the budget notification to when the threshold is exceeded.",
            nullable=True,
        )
        _element.contact_roles = AAZListArg(
            options=["contact-roles"],
            help="Contact roles to send the budget notification to when the threshold is exceeded.",
            nullable=True,
        )
        _element.enabled = AAZBoolArg(
            options=["enabled"],
            help="The notification is enabled or not.",
        )
        _element.operator = AAZStrArg(
            options=["operator"],
            help="The comparison operator.",
            enum={"EqualTo": "EqualTo", "GreaterThan": "GreaterThan", "GreaterThanOrEqualTo": "GreaterThanOrEqualTo"},
        )
        _element.threshold = AAZFloatArg(
            options=["threshold"],
            help="Threshold value associated with a notification. Notification is sent when the cost exceeded the threshold. It is always percent and has to be between 0 and 1000.",
        )

        contact_emails = cls._args_schema.notifications.Element.contact_emails
        contact_emails.Element = AAZStrArg(
            nullable=True,
        )

        contact_groups = cls._args_schema.notifications.Element.contact_groups
        contact_groups.Element = AAZStrArg(
            nullable=True,
        )

        contact_roles = cls._args_schema.notifications.Element.contact_roles
        contact_roles.Element = AAZStrArg(
            nullable=True,
        )

        time_period = cls._args_schema.time_period
        time_period.end_date = AAZDateTimeArg(
            options=["end-date"],
            help="The end date for the budget. If not provided, we default this to 10 years from the start date.",
            nullable=True,
        )
        time_period.start_date = AAZDateTimeArg(
            options=["start-date"],
            help="The start date for the budget.",
        )
        return cls._args_schema

    def _execute_operations(self):
        self.pre_operations()
        self.BudgetsGet(ctx=self.ctx)()
        self.pre_instance_update(self.ctx.vars.instance)
        self.InstanceUpdateByJson(ctx=self.ctx)()
        self.InstanceUpdateByGeneric(ctx=self.ctx)()
        self.post_instance_update(self.ctx.vars.instance)
        self.BudgetsCreateOrUpdate(ctx=self.ctx)()
        self.post_operations()

    @register_callback
    def pre_operations(self):
        pass

    @register_callback
    def post_operations(self):
        pass

    @register_callback
    def pre_instance_update(self, instance):
        pass

    @register_callback
    def post_instance_update(self, instance):
        pass

    def _output(self, *args, **kwargs):
        result = self.deserialize_output(self.ctx.vars.instance, client_flatten=True)
        return result

    class BudgetsGet(AAZHttpOperation):
        CLIENT_TYPE = "MgmtClient"

        def __call__(self, *args, **kwargs):
            request = self.make_request()
            session = self.client.send_request(request=request, stream=False, **kwargs)
            if session.http_response.status_code in [200]:
                return self.on_200(session)

            return self.on_error(session.http_response)

        @property
        def url(self):
            return self.client.format_url(
                "/subscriptions/{subscriptionId}/providers/Microsoft.Consumption/budgets/{budgetName}",
                **self.url_parameters
            )

        @property
        def method(self):
            return "GET"

        @property
        def error_format(self):
            return "ODataV4Format"

        @property
        def url_parameters(self):
            parameters = {
                **self.serialize_url_param(
                    "budgetName", self.ctx.args.budget_name,
                    required=True,
                ),
                **self.serialize_url_param(
                    "subscriptionId", self.ctx.subscription_id,
                    required=True,
                ),
            }
            return parameters

        @property
        def query_parameters(self):
            parameters = {
                **self.serialize_query_param(
                    "api-version", "2023-05-01",
                    required=True,
                ),
            }
            return parameters

        @property
        def header_parameters(self):
            parameters = {
                **self.serialize_header_param(
                    "Accept", "application/json",
                ),
            }
            return parameters

        def on_200(self, session):
            data = self.deserialize_http_content(session)
            self.ctx.set_var(
                "instance",
                data,
                schema_builder=self._build_schema_on_200
            )

        _schema_on_200 = None

        @classmethod
        def _build_schema_on_200(cls):
            if cls._schema_on_200 is not None:
                return cls._schema_on_200

            cls._schema_on_200 = AAZObjectType()
            _UpdateHelper._build_schema_budget_read(cls._schema_on_200)

            return cls._schema_on_200

    class BudgetsCreateOrUpdate(AAZHttpOperation):
        CLIENT_TYPE = "MgmtClient"

        def __call__(self, *args, **kwargs):
            request = self.make_request()
            session = self.client.send_request(request=request, stream=False, **kwargs)
            if session.http_response.status_code in [200, 201]:
                return self.on_200_201(session)

            return self.on_error(session.http_response)

        @property
        def url(self):
            return self.client.format_url(
                "/subscriptions/{subscriptionId}/providers/Microsoft.Consumption/budgets/{budgetName}",
                **self.url_parameters
            )

        @property
        def method(self):
            return "PUT"

        @property
        def error_format(self):
            return "ODataV4Format"

        @property
        def url_parameters(self):
            parameters = {
                **self.serialize_url_param(
                    "budgetName", self.ctx.args.budget_name,
                    required=True,
                ),
                **self.serialize_url_param(
                    "subscriptionId", self.ctx.subscription_id,
                    required=True,
                ),
            }
            return parameters

        @property
        def query_parameters(self):
            parameters = {
                **self.serialize_query_param(
                    "api-version", "2023-05-01",
                    required=True,
                ),
            }
            return parameters

        @property
        def header_parameters(self):
            parameters = {
                **self.serialize_header_param(
                    "Content-Type", "application/json",
                ),
                **self.serialize_header_param(
                    "Accept", "application/json",
                ),
            }
            return parameters

        @property
        def content(self):
            _content_value, _builder = self.new_content_builder(
                self.ctx.args,
                value=self.ctx.vars.instance,
            )

            return self.serialize_content(_content_value)

        def on_200_201(self, session):
            data = self.deserialize_http_content(session)
            self.ctx.set_var(
                "instance",
                data,
                schema_builder=self._build_schema_on_200_201
            )

        _schema_on_200_201 = None

        @classmethod
        def _build_schema_on_200_201(cls):
            if cls._schema_on_200_201 is not None:
                return cls._schema_on_200_201

            cls._schema_on_200_201 = AAZObjectType()
            _UpdateHelper._build_schema_budget_read(cls._schema_on_200_201)

            return cls._schema_on_200_201

    class InstanceUpdateByJson(AAZJsonInstanceUpdateOperation):

        def __call__(self, *args, **kwargs):
            self._update_instance(self.ctx.vars.instance)

        def _update_instance(self, instance):
            _instance_value, _builder = self.new_content_builder(
                self.ctx.args,
                value=instance,
                typ=AAZObjectType
            )
            _builder.set_prop("eTag", AAZStrType, ".e_tag")
            _builder.set_prop("properties", AAZObjectType, typ_kwargs={"flags": {"client_flatten": True}})

            properties = _builder.get(".properties")
            if properties is not None:
                properties.set_prop("amount", AAZFloatType, ".amount", typ_kwargs={"flags": {"required": True}})
                properties.set_prop("category", AAZStrType, ".category", typ_kwargs={"flags": {"required": True}})
                properties.set_prop("filters", AAZObjectType, ".filters")
                properties.set_prop("notifications", AAZDictType, ".notifications")
                properties.set_prop("timeGrain", AAZStrType, ".time_grain", typ_kwargs={"flags": {"required": True}})
                properties.set_prop("timePeriod", AAZObjectType, ".time_period", typ_kwargs={"flags": {"required": True}})

            filters = _builder.get(".properties.filters")
            if filters is not None:
                filters.set_prop("meters", AAZListType, ".meters")
                filters.set_prop("resourceGroups", AAZListType, ".resource_groups")
                filters.set_prop("resources", AAZListType, ".resources")

            meters = _builder.get(".properties.filters.meters")
            if meters is not None:
                meters.set_elements(AAZStrType, ".")

            resource_groups = _builder.get(".properties.filters.resourceGroups")
            if resource_groups is not None:
                resource_groups.set_elements(AAZStrType, ".")

            resources = _builder.get(".properties.filters.resources")
            if resources is not None:
                resources.set_elements(AAZStrType, ".")

            notifications = _builder.get(".properties.notifications")
            if notifications is not None:
                notifications.set_elements(AAZObjectType, ".")

            _elements = _builder.get(".properties.notifications{}")
            if _elements is not None:
                _elements.set_prop("contactEmails", AAZListType, ".contact_emails", typ_kwargs={"flags": {"required": True}})
                _elements.set_prop("contactGroups", AAZListType, ".contact_groups")
                _elements.set_prop("contactRoles", AAZListType, ".contact_roles")
                _elements.set_prop("enabled", AAZBoolType, ".enabled", typ_kwargs={"flags": {"required": True}})
                _elements.set_prop("operator", AAZStrType, ".operator", typ_kwargs={"flags": {"required": True}})
                _elements.set_prop("threshold", AAZFloatType, ".threshold", typ_kwargs={"flags": {"required": True}})

            contact_emails = _builder.get(".properties.notifications{}.contactEmails")
            if contact_emails is not None:
                contact_emails.set_elements(AAZStrType, ".")

            contact_groups = _builder.get(".properties.notifications{}.contactGroups")
            if contact_groups is not None:
                contact_groups.set_elements(AAZStrType, ".")

            contact_roles = _builder.get(".properties.notifications{}.contactRoles")
            if contact_roles is not None:
                contact_roles.set_elements(AAZStrType, ".")

            time_period = _builder.get(".properties.timePeriod")
            if time_period is not None:
                time_period.set_prop("endDate", AAZStrType, ".end_date")
                time_period.set_prop("startDate", AAZStrType, ".start_date", typ_kwargs={"flags": {"required": True}})

            return _instance_value

    class InstanceUpdateByGeneric(AAZGenericInstanceUpdateOperation):

        def __call__(self, *args, **kwargs):
            self._update_instance_by_generic(
                self.ctx.vars.instance,
                self.ctx.generic_update_args
            )


class _UpdateHelper:
    """Helper class for Update"""

    _schema_budget_read = None

    @classmethod
    def _build_schema_budget_read(cls, _schema):
        if cls._schema_budget_read is not None:
            _schema.e_tag = cls._schema_budget_read.e_tag
            _schema.id = cls._schema_budget_read.id
            _schema.name = cls._schema_budget_read.name
            _schema.properties = cls._schema_budget_read.properties
            _schema.type = cls._schema_budget_read.type
            return

        cls._schema_budget_read = _schema_budget_read = AAZObjectType()

        budget_read = _schema_budget_read
        budget_read.e_tag = AAZStrType(
            serialized_name="eTag",
        )
        budget_read.id = AAZStrType(
            flags={"read_only": True},
        )
        budget_read.name = AAZStrType(
            flags={"read_only": True},
        )
        budget_read.properties = AAZObjectType(
            flags={"client_flatten": True},
        )
        budget_read.type = AAZStrType(
            flags={"read_only": True},
        )

        properties = _schema_budget_read.properties
        properties.amount = AAZFloatType(
            flags={"required": True},
        )
        properties.category = AAZStrType(
            flags={"required": True},
        )
        properties.current_spend = AAZObjectType(
            serialized_name="currentSpend",
        )
        properties.filters = AAZObjectType()
        properties.notifications = AAZDictType()
        properties.time_grain = AAZStrType(
            serialized_name="timeGrain",
            flags={"required": True},
        )
        properties.time_period = AAZObjectType(
            serialized_name="timePeriod",
            flags={"required": True},
        )

        current_spend = _schema_budget_read.properties.current_spend
        current_spend.amount = AAZFloatType(
            flags={"read_only": True},
        )
        current_spend.unit = AAZStrType(
            flags={"read_only": True},
        )

        filters = _schema_budget_read.properties.filters
        filters.meters = AAZListType()
        filters.resource_groups = AAZListType(
            serialized_name="resourceGroups",
        )
        filters.resources = AAZListType()

        meters = _schema_budget_read.properties.filters.meters
        meters.Element = AAZStrType()

        resource_groups = _schema_budget_read.properties.filters.resource_groups
        resource_groups.Element = AAZStrType()

        resources = _schema_budget_read.properties.filters.resources
        resources.Element = AAZStrType()

        notifications = _schema_budget_read.properties.notifications
        notifications.Element = AAZObjectType()

        _element = _schema_budget_read.properties.notifications.Element
        _element.contact_emails = AAZListType(
            serialized_name="contactEmails",
            flags={"required": True},
        )
        _element.contact_groups = AAZListType(
            serialized_name="contactGroups",
        )
        _element.contact_roles = AAZListType(
            serialized_name="contactRoles",
        )
        _element.enabled = AAZBoolType(
            flags={"required": True},
        )
        _element.operator = AAZStrType(
            flags={"required": True},
        )
        _element.threshold = AAZFloatType(
            flags={"required": True},
        )

        contact_emails = _schema_budget_read.properties.notifications.Element.contact_emails
        contact_emails.Element = AAZStrType()

        contact_groups = _schema_budget_read.properties.notifications.Element.contact_groups
        contact_groups.Element = AAZStrType()

        contact_roles = _schema_budget_read.properties.notifications.Element.contact_roles
        contact_roles.Element = AAZStrType()

        time_period = _schema_budget_read.properties.time_period
        time_period.end_date = AAZStrType(
            serialized_name="endDate",
        )
        time_period.start_date = AAZStrType(
            serialized_name="startDate",
            flags={"required": True},
        )

        _schema.e_tag = cls._schema_budget_read.e_tag
        _schema.id = cls._schema_budget_read.id
        _schema.name = cls._schema_budget_read.name
        _schema.properties = cls._schema_budget_read.properties
        _schema.type = cls._schema_budget_read.type


__all__ = ["Update"]
