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

# Copyright: (c) 2019, Michael Pechner <mikey@mikey.com>
# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)

DOCUMENTATION = r"""
---
module: ecs_tag
version_added: 1.0.0
short_description: create and remove tags on Amazon ECS resources
description:
  - Creates and removes tags for Amazon ECS resources.
  - Resources are referenced by their cluster name.
author:
  - Michael Pechner (@mpechner)
options:
  cluster_name:
    description:
      - The name of the cluster whose resources we are tagging.
    required: true
    type: str
  resource:
    description:
      - The ECS resource name.
      - Required unless I(resource_type=cluster).
    type: str
  resource_type:
    description:
      - The type of resource.
    default: cluster
    choices: ['cluster', 'task', 'service', 'task_definition', 'container']
    type: str
  state:
    description:
      - Whether the tags should be present or absent on the resource.
    default: present
    choices: ['present', 'absent']
    type: str
  tags:
    description:
      - A dictionary of tags to add or remove from the resource.
      - If the value provided for a tag is null and I(state=absent), the tag will be removed regardless of its current value.
    type: dict
    aliases: ['resource_tags']
  purge_tags:
    description:
      - Whether unspecified tags should be removed from the resource.
      - Note that when combined with I(state=absent), specified tags with non-matching values are not purged.
    type: bool
    default: false
extends_documentation_fragment:
  - amazon.aws.common.modules
  - amazon.aws.region.modules
  - amazon.aws.boto3
"""

EXAMPLES = r"""
- name: Ensure tags are present on a resource
  community.aws.ecs_tag:
    cluster_name: mycluster
    resource_type: cluster
    state: present
    tags:
      Name: ubervol
      env: prod

- name: Remove the Env tag
  community.aws.ecs_tag:
    cluster_name: mycluster
    resource_type: cluster
    tags:
      Env:
    state: absent

- name: Remove the Env tag if it's currently 'development'
  community.aws.ecs_tag:
    cluster_name: mycluster
    resource_type: cluster
    tags:
      Env: development
    state: absent

- name: Remove all tags except for Name from a cluster
  community.aws.ecs_tag:
    cluster_name: mycluster
    resource_type: cluster
    tags:
        Name: foo
    state: absent
    purge_tags: true
"""

RETURN = r"""
tags:
  description: A dict containing the tags on the resource
  returned: always
  type: dict
added_tags:
  description: A dict of tags that were added to the resource
  returned: If tags were added
  type: dict
removed_tags:
  description: A dict of tags that were removed from the resource
  returned: If tags were removed
  type: dict
"""

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

from ansible_collections.amazon.aws.plugins.module_utils.tagging import ansible_dict_to_boto3_tag_list
from ansible_collections.amazon.aws.plugins.module_utils.tagging import boto3_tag_list_to_ansible_dict
from ansible_collections.amazon.aws.plugins.module_utils.tagging import compare_aws_tags

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


def get_tags(ecs, module, resource):
    try:
        return boto3_tag_list_to_ansible_dict(ecs.list_tags_for_resource(resourceArn=resource)["tags"])
    except (BotoCoreError, ClientError) as e:
        module.fail_json_aws(e, msg=f"Failed to fetch tags for resource {resource}")


def get_arn(ecs, module, cluster_name, resource_type, resource):
    try:
        if resource_type == "cluster":
            description = ecs.describe_clusters(clusters=[resource])
            resource_arn = description["clusters"][0]["clusterArn"]
        elif resource_type == "task":
            description = ecs.describe_tasks(cluster=cluster_name, tasks=[resource])
            resource_arn = description["tasks"][0]["taskArn"]
        elif resource_type == "service":
            description = ecs.describe_services(cluster=cluster_name, services=[resource])
            resource_arn = description["services"][0]["serviceArn"]
        elif resource_type == "task_definition":
            description = ecs.describe_task_definition(taskDefinition=resource)
            resource_arn = description["taskDefinition"]["taskDefinitionArn"]
        elif resource_type == "container":
            description = ecs.describe_container_instances(clusters=[resource])
            resource_arn = description["containerInstances"][0]["containerInstanceArn"]
    except (IndexError, KeyError):
        module.fail_json(msg=f"Failed to find {resource_type} {resource}")
    except (BotoCoreError, ClientError) as e:
        module.fail_json_aws(e, msg=f"Failed to find {resource_type} {resource}")

    return resource_arn


def main():
    argument_spec = dict(
        cluster_name=dict(required=True),
        resource=dict(required=False),
        tags=dict(type="dict", aliases=["resource_tags"]),
        purge_tags=dict(type="bool", default=False),
        state=dict(default="present", choices=["present", "absent"]),
        resource_type=dict(default="cluster", choices=["cluster", "task", "service", "task_definition", "container"]),
    )
    required_if = [("state", "present", ["tags"]), ("state", "absent", ["tags"])]

    module = AnsibleAWSModule(argument_spec=argument_spec, required_if=required_if, supports_check_mode=True)

    resource_type = module.params["resource_type"]
    cluster_name = module.params["cluster_name"]
    if resource_type == "cluster":
        resource = cluster_name
    else:
        resource = module.params["resource"]
    tags = module.params["tags"]
    state = module.params["state"]
    purge_tags = module.params["purge_tags"]

    result = {"changed": False}

    ecs = module.client("ecs")

    resource_arn = get_arn(ecs, module, cluster_name, resource_type, resource)

    current_tags = get_tags(ecs, module, resource_arn)

    add_tags, remove = compare_aws_tags(current_tags, tags, purge_tags=purge_tags)

    remove_tags = {}
    if state == "absent":
        for key in tags:
            if key in current_tags and (tags[key] is None or current_tags[key] == tags[key]):
                remove_tags[key] = current_tags[key]

    for key in remove:
        remove_tags[key] = current_tags[key]

    if remove_tags:
        result["changed"] = True
        result["removed_tags"] = remove_tags
        if not module.check_mode:
            try:
                ecs.untag_resource(resourceArn=resource_arn, tagKeys=list(remove_tags.keys()))
            except (BotoCoreError, ClientError) as e:
                module.fail_json_aws(e, msg=f"Failed to remove tags {remove_tags} from resource {resource}")

    if state == "present" and add_tags:
        result["changed"] = True
        result["added_tags"] = add_tags
        current_tags.update(add_tags)
        if not module.check_mode:
            try:
                tags = ansible_dict_to_boto3_tag_list(add_tags, tag_name_key_name="key", tag_value_key_name="value")
                ecs.tag_resource(resourceArn=resource_arn, tags=tags)
            except (BotoCoreError, ClientError) as e:
                module.fail_json_aws(e, msg=f"Failed to set tags {add_tags} on resource {resource}")

    result["tags"] = get_tags(ecs, module, resource_arn)
    module.exit_json(**result)


if __name__ == "__main__":
    main()
