#!/usr/bin/python
# -*- coding: utf-8 -*-

# Copyright: Ansible Project
# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)

DOCUMENTATION = r"""
module: ec2_transit_gateway
short_description: Create and delete AWS Transit Gateways
version_added: 1.0.0
description:
  - Creates AWS Transit Gateways.
  - Deletes AWS Transit Gateways.
  - Updates tags on existing transit gateways.
options:
  asn:
    description:
      - A private Autonomous System Number (ASN) for the Amazon side of a BGP session.
      - The range is 64512 to 65534 for 16-bit ASNs and 4200000000 to 4294967294 for 32-bit ASNs.
    type: int
  auto_associate:
    description:
      - Enable or disable automatic association with the default association route table.
    default: true
    type: bool
  auto_attach:
    description:
      - Enable or disable automatic acceptance of attachment requests.
    default: false
    type: bool
  auto_propagate:
    description:
      - Enable or disable automatic propagation of routes to the default propagation route table.
    default: true
    type: bool
  description:
     description:
       - The description of the transit gateway.
     type: str
  dns_support:
    description:
      - Whether to enable AWS DNS support.
    default: true
    type: bool
  state:
    description:
      - C(present) to ensure resource is created.
      - C(absent) to remove resource.
    default: present
    choices: [ "present", "absent"]
    type: str
  transit_gateway_id:
    description:
      - The ID of the transit gateway.
    type: str
  vpn_ecmp_support:
    description:
      - Enable or disable Equal Cost Multipath Protocol support.
    default: true
    type: bool
  wait:
    description:
      - Whether to wait for status
    default: true
    type: bool
  wait_timeout:
    description:
      - number of seconds to wait for status
    default: 300
    type: int

author:
  - "Bob Boldin (@BobBoldin)"
extends_documentation_fragment:
  - amazon.aws.common.modules
  - amazon.aws.region.modules
  - amazon.aws.boto3
  - amazon.aws.tags
"""

EXAMPLES = r"""
- name: Create a new transit gateway using defaults
  community.aws.ec2_transit_gateway:
    state: present
    region: us-east-1
    description: personal-testing
  register: created_tgw

- name: Create a new transit gateway with options
  community.aws.ec2_transit_gateway:
    asn: 64514
    auto_associate: false
    auto_propagate: false
    dns_support: true
    description: "nonprod transit gateway"
    purge_tags: false
    state: present
    region: us-east-1
    tags:
      Name: nonprod transit gateway
      status: testing

- name: Remove a transit gateway by description
  community.aws.ec2_transit_gateway:
    state: absent
    region: us-east-1
    description: personal-testing

- name: Remove a transit gateway by id
  community.aws.ec2_transit_gateway:
    state: absent
    region: ap-southeast-2
    transit_gateway_id: tgw-3a9aa123
  register: deleted_tgw
"""

RETURN = r"""
transit_gateway:
  description: The attributes of the transit gateway.
  type: complex
  returned: I(state=present)
  contains:
    creation_time:
      description: The creation time of the transit gateway.
      returned: always
      type: str
      sample: "2019-03-06T17:13:51+00:00"
    description:
      description: The description of the transit gateway.
      returned: always
      type: str
      sample: my test tgw
    options:
      description: The options attributes of the transit gateway
      returned: always
      type: complex
      contains:
        amazon_side_asn:
          description:
            - A private Autonomous System Number (ASN) for the Amazon side of a BGP session.
              The range is 64512 to 65534 for 16-bit ASNs and 4200000000 to 4294967294 for 32-bit ASNs.
          returned: always
          type: str
          sample: 64512
        auto_accept_shared_attachements:
          description: Indicates whether attachment requests are automatically accepted.
          returned: always
          type: str
          sample: disable
        default_route_table_association:
          description:
           - Indicates  whether resource attachments are automatically
              associated with the default association route table.
          returned: always
          type: str
          sample: enable
        association_default_route_table_id:
          description: The ID of the default association route table.
          returned: Iwhen exists
          type: str
          sample: tgw-rtb-abc123444
        default_route_table_propagation:
          description:
           - Indicates  whether  resource  attachments   automatically
             propagate routes to the default propagation route table.
          returned: always
          type: str
          sample: disable
        propagation_default_route_table_id:
          description: The ID of the default propagation route table.
          returned: when exists
          type: str
          sample: tgw-rtb-def456777
        vpn_ecmp_support:
          description: Indicates  whether  Equal Cost Multipath Protocol support is enabled.
          returned: always
          type: str
          sample: enable
        dns_support:
          description: Indicates whether DNS support is enabled.
          returned: always
          type: str
          sample: enable
    owner_id:
      description: The account that owns the transit gateway.
      returned: always
      type: str
      sample: '123456789012'
    state:
      description: The state of the transit gateway.
      returned: always
      type: str
      sample: pending
    tags:
      description: A dictionary of resource tags
      returned: always
      type: dict
      sample:
        tags:
          Name: nonprod_tgw
    transit_gateway_arn:
      description: The ID of the transit_gateway.
      returned: always
      type: str
      sample: tgw-3a9aa123
    transit_gateway_id:
      description: The ID of the transit_gateway.
      returned: always
      type: str
      sample: tgw-3a9aa123
"""

from time import sleep
from time import time

try:
    from botocore.exceptions import BotoCoreError
    from botocore.exceptions import ClientError
except ImportError:
    pass  # handled by imported AnsibleAWSModule

from ansible.module_utils.common.dict_transformations import camel_dict_to_snake_dict

from ansible_collections.amazon.aws.plugins.module_utils.ec2 import ensure_ec2_tags
from ansible_collections.amazon.aws.plugins.module_utils.retries import AWSRetry
from ansible_collections.amazon.aws.plugins.module_utils.tagging import boto3_tag_list_to_ansible_dict
from ansible_collections.amazon.aws.plugins.module_utils.transformation import ansible_dict_to_boto3_filter_list

from ansible_collections.community.aws.plugins.module_utils.modules import AnsibleCommunityAWSModule as AnsibleAWSModule


class AnsibleEc2Tgw(object):
    def __init__(self, module, results):
        self._module = module
        self._results = results
        retry_decorator = AWSRetry.jittered_backoff(
            catch_extra_error_codes=["IncorrectState"],
        )
        connection = module.client("ec2", retry_decorator=retry_decorator)
        self._connection = connection
        self._check_mode = self._module.check_mode

    def process(self):
        """Process the request based on state parameter .
        state = present will search for an existing tgw based and return the object data.
          if no object is found it will be created

        state = absent will attempt to remove the tgw however will fail if it still has
          attachments or associations
        """
        description = self._module.params.get("description")
        state = self._module.params.get("state", "present")
        tgw_id = self._module.params.get("transit_gateway_id")

        if state == "present":
            self.ensure_tgw_present(tgw_id, description)
        elif state == "absent":
            self.ensure_tgw_absent(tgw_id, description)

    def wait_for_status(self, wait_timeout, tgw_id, status, skip_deleted=True):
        """
        Wait for the Transit Gateway to reach  the specified status
        :param wait_timeout: Number of seconds to wait, until this timeout is reached.
        :param tgw_id: The Amazon nat id.
        :param status: The status to wait for.
                examples. status=available, status=deleted
        :param skip_deleted: ignore deleted transit gateways
        :return dict: transit gateway object
        """
        polling_increment_secs = 5
        wait_timeout = time() + wait_timeout
        status_achieved = False
        transit_gateway = dict()

        while wait_timeout > time():
            try:
                transit_gateway = self.get_matching_tgw(tgw_id=tgw_id, skip_deleted=skip_deleted)

                if transit_gateway:
                    if self._check_mode:
                        transit_gateway["state"] = status

                    if transit_gateway.get("state") == status:
                        status_achieved = True
                        break

                    elif transit_gateway.get("state") == "failed":
                        break

                else:
                    sleep(polling_increment_secs)

            except ClientError as e:
                self._module.fail_json_aws(e)

        if not status_achieved:
            self._module.fail_json(msg="Wait time out reached, while waiting for results")

        return transit_gateway

    def get_matching_tgw(self, tgw_id, description=None, skip_deleted=True):
        """search for  an existing tgw by either tgw_id or description
        :param tgw_id:  The AWS id of the transit gateway
        :param description:  The description of the transit gateway.
        :param skip_deleted: ignore deleted transit gateways
        :return dict: transit gateway object
        """
        filters = []
        if tgw_id:
            filters = ansible_dict_to_boto3_filter_list({"transit-gateway-id": tgw_id})

        try:
            response = AWSRetry.exponential_backoff()(self._connection.describe_transit_gateways)(Filters=filters)
        except (ClientError, BotoCoreError) as e:
            self._module.fail_json_aws(e)

        tgw = None
        tgws = []

        if len(response.get("TransitGateways", [])) == 1 and tgw_id:
            if (response["TransitGateways"][0]["State"] != "deleted") or not skip_deleted:
                tgws.extend(response["TransitGateways"])

        for gateway in response.get("TransitGateways", []):
            if description == gateway["Description"] and gateway["State"] != "deleted":
                tgws.append(gateway)

        if len(tgws) > 1:
            self._module.fail_json(
                msg=f"EC2 returned more than one transit Gateway for description {description}, aborting"
            )
        elif tgws:
            tgw = camel_dict_to_snake_dict(tgws[0], ignore_list=["Tags"])
            tgw["tags"] = boto3_tag_list_to_ansible_dict(tgws[0]["Tags"])

        return tgw

    @staticmethod
    def enable_option_flag(flag):
        disabled = "disable"
        enabled = "enable"
        if flag:
            return enabled
        return disabled

    def create_tgw(self, description):
        """
        Create a transit gateway and optionally wait for status to become available.

        :param description: The description of the transit gateway.
        :return dict: transit gateway object
        """
        options = dict()
        wait = self._module.params.get("wait")
        wait_timeout = self._module.params.get("wait_timeout")

        if self._module.params.get("asn"):
            options["AmazonSideAsn"] = self._module.params.get("asn")

        options["AutoAcceptSharedAttachments"] = self.enable_option_flag(self._module.params.get("auto_attach"))
        options["DefaultRouteTableAssociation"] = self.enable_option_flag(self._module.params.get("auto_associate"))
        options["DefaultRouteTablePropagation"] = self.enable_option_flag(self._module.params.get("auto_propagate"))
        options["VpnEcmpSupport"] = self.enable_option_flag(self._module.params.get("vpn_ecmp_support"))
        options["DnsSupport"] = self.enable_option_flag(self._module.params.get("dns_support"))

        try:
            response = self._connection.create_transit_gateway(Description=description, Options=options)
        except (ClientError, BotoCoreError) as e:
            self._module.fail_json_aws(e)

        tgw_id = response["TransitGateway"]["TransitGatewayId"]

        if wait:
            result = self.wait_for_status(wait_timeout=wait_timeout, tgw_id=tgw_id, status="available")
        else:
            result = self.get_matching_tgw(tgw_id=tgw_id)

        self._results["msg"] = f" Transit gateway {result['transit_gateway_id']} created"

        return result

    def delete_tgw(self, tgw_id):
        """
        De;lete the transit gateway and optionally wait for status to become deleted

        :param tgw_id: The id of the transit gateway
        :return dict: transit gateway object
        """
        wait = self._module.params.get("wait")
        wait_timeout = self._module.params.get("wait_timeout")

        try:
            response = self._connection.delete_transit_gateway(TransitGatewayId=tgw_id)
        except (ClientError, BotoCoreError) as e:
            self._module.fail_json_aws(e)

        if wait:
            result = self.wait_for_status(
                wait_timeout=wait_timeout, tgw_id=tgw_id, status="deleted", skip_deleted=False
            )
        else:
            result = self.get_matching_tgw(tgw_id=tgw_id, skip_deleted=False)

        self._results["msg"] = f" Transit gateway {tgw_id} deleted"

        return result

    def ensure_tgw_present(self, tgw_id=None, description=None):
        """
        Will create a tgw if no match to the tgw_id or description are found
        Will update the tgw tags if matching one found but tags are not synced

        :param tgw_id:  The AWS id of the transit gateway
        :param description:  The description of the transit gateway.
        :return dict: transit gateway object
        """
        tgw = self.get_matching_tgw(tgw_id, description)

        if tgw is None:
            if self._check_mode:
                self._results["changed"] = True
                self._results["transit_gateway_id"] = None
                return self._results

            try:
                if not description:
                    self._module.fail_json(msg="Failed to create Transit Gateway: description argument required")
                tgw = self.create_tgw(description)
                self._results["changed"] = True
            except (BotoCoreError, ClientError) as e:
                self._module.fail_json_aws(e, msg="Unable to create Transit Gateway")

        self._results["changed"] |= ensure_ec2_tags(
            self._connection,
            self._module,
            tgw["transit_gateway_id"],
            tags=self._module.params.get("tags"),
            purge_tags=self._module.params.get("purge_tags"),
        )

        self._results["transit_gateway"] = self.get_matching_tgw(tgw_id=tgw["transit_gateway_id"])

        return self._results

    def ensure_tgw_absent(self, tgw_id=None, description=None):
        """
        Will delete the tgw if a single tgw is found not yet in deleted status

        :param tgw_id:  The AWS id of the transit gateway
        :param description:  The description of the transit gateway.
        :return doct: transit gateway object
        """
        self._results["transit_gateway_id"] = None
        tgw = self.get_matching_tgw(tgw_id, description)

        if tgw is not None:
            if self._check_mode:
                self._results["changed"] = True
                return self._results

            try:
                tgw = self.delete_tgw(tgw_id=tgw["transit_gateway_id"])
                self._results["changed"] = True
                self._results["transit_gateway"] = self.get_matching_tgw(
                    tgw_id=tgw["transit_gateway_id"], skip_deleted=False
                )
            except (BotoCoreError, ClientError) as e:
                self._module.fail_json_aws(e, msg="Unable to delete Transit Gateway")

        return self._results


def setup_module_object():
    """
    merge argument spec and create Ansible module object
    :return: Ansible module object
    """

    argument_spec = dict(
        asn=dict(type="int"),
        auto_associate=dict(type="bool", default=True),
        auto_attach=dict(type="bool", default=False),
        auto_propagate=dict(type="bool", default=True),
        description=dict(type="str"),
        dns_support=dict(type="bool", default=True),
        purge_tags=dict(type="bool", default=True),
        state=dict(default="present", choices=["present", "absent"]),
        tags=dict(type="dict", aliases=["resource_tags"]),
        transit_gateway_id=dict(type="str"),
        vpn_ecmp_support=dict(type="bool", default=True),
        wait=dict(type="bool", default=True),
        wait_timeout=dict(type="int", default=300),
    )

    module = AnsibleAWSModule(
        argument_spec=argument_spec,
        required_one_of=[("description", "transit_gateway_id")],
        supports_check_mode=True,
    )

    return module


def main():
    module = setup_module_object()

    results = dict(changed=False)

    tgw_manager = AnsibleEc2Tgw(module=module, results=results)
    tgw_manager.process()

    module.exit_json(**results)


if __name__ == "__main__":
    main()
