import boto3

from moto import mock_aws

from .test_ecs import setup_ecs_cluster_with_ec2_instance


@mock_aws
def test_describe_tasks_include_tags():
    client = boto3.client("ecs", region_name="us-east-1")
    test_cluster_name = "test_ecs_cluster"
    setup_ecs_cluster_with_ec2_instance(client, test_cluster_name)

    task_tags = [{"key": "Name", "value": "test_ecs_task"}]
    tasks_arns = [
        task["taskArn"]
        for task in client.run_task(
            cluster="test_ecs_cluster",
            taskDefinition="test_ecs_task",
            count=2,
            tags=task_tags,
        )["tasks"]
    ]
    response = client.describe_tasks(
        cluster="test_ecs_cluster", tasks=tasks_arns, include=["TAGS"]
    )

    assert len(response["tasks"]) == 2
    assert {response["tasks"][0]["taskArn"], response["tasks"][1]["taskArn"]} == set(
        tasks_arns
    )
    assert response["tasks"][0]["tags"] == task_tags

    # Test we can pass task ids instead of ARNs
    response = client.describe_tasks(
        cluster="test_ecs_cluster", tasks=[tasks_arns[0].split("/")[-1]]
    )
    assert len(response["tasks"]) == 1

    tags = client.list_tags_for_resource(resourceArn=tasks_arns[0])["tags"]
    assert tags == task_tags


@mock_aws
def test_add_tags_to_task():
    client = boto3.client("ecs", region_name="us-east-1")
    test_cluster_name = "test_ecs_cluster"
    setup_ecs_cluster_with_ec2_instance(client, test_cluster_name)

    task_tags = [{"key": "k1", "value": "v1"}]
    task_arn = client.run_task(
        cluster="test_ecs_cluster",
        taskDefinition="test_ecs_task",
        count=1,
        tags=task_tags,
    )["tasks"][0]["taskArn"]

    client.tag_resource(resourceArn=task_arn, tags=[{"key": "k2", "value": "v2"}])

    tags = client.describe_tasks(
        cluster="test_ecs_cluster", tasks=[task_arn], include=["TAGS"]
    )["tasks"][0]["tags"]
    assert len(tags) == 2
    assert {"key": "k1", "value": "v1"} in tags
    assert {"key": "k2", "value": "v2"} in tags

    client.untag_resource(resourceArn=task_arn, tagKeys=["k2"])

    resp = client.list_tags_for_resource(resourceArn=task_arn)
    assert resp["tags"] == [{"key": "k1", "value": "v1"}]
