import boto3
import json
import botocore
import os
from botocore.config import Config


def get_available_region_list():
    ec2 = boto3.client('ec2')
    response = ec2.describe_regions()
    region_list = []
    for region in response['Regions']:
        region_list.append(region['RegionName'])
    return region_list


def get_aws_resources_elb_v2(region):
    config = Config(
        retries={
            'max_attempts': 10,
            'mode': 'adaptive'
        }
    )
    elb_output = []
    elbclient = boto3.client('elbv2', region_name=region, config=config)
    response = elbclient.describe_load_balancers()
    for elb in response['LoadBalancers']:
        tags = elbclient.describe_tags(
            ResourceArns=[elb['LoadBalancerArn']])
        elb.update({'Tags': tags['TagDescriptions']})
        elb_output.append(elb)
    return elb_output


def tag_ec2(ec2, region, cluster_list):
    ec2client = boto3.client('ec2', region_name=region)
    for reservation in ec2['Reservations']:
        for instance in reservation['Instances']:
            cluster_name = ''
            is_devtier = False
            if 'Tags' in instance:
                for tag in instance['Tags']:
                    if 'kubernetes.io/cluster/' in tag['Key']:
                        cluster_name = tag['Key'].replace(
                            'kubernetes.io/cluster/', '')
                        if cluster_name in cluster_list:
                            is_devtier = True
                            env_name = cluster_name.split('-')[0]
                        else:
                            # not the devtier resources
                            break
                    if 'servicetype' in tag['Key']:
                        # has been labeled
                        break
            if is_devtier:
                ec2client.create_tags(Resources=[instance['InstanceId']], Tags=[{'Key': 'component', 'Value': 'devtier-control-plane'}, {'Key': 'servicetype', 'Value': 'devtier'}, {
                    'Key': 'environment', 'Value': env_name}, {'Key': 'tenant', 'Value': 'admin'}, {'Key': 'usedby', 'Value': "%s" % cluster_name}])


def tag_ebs(ebs, region):
    ec2client = boto3.client('ec2', region_name=region)
    for volume in ebs['Volumes']:
        is_devtier = False
        component_name = ''
        if 'Tags' in volume:
            for tag in volume['Tags']:
                if 'kubernetes.io/created-for/pvc/namespace' in tag['Key']:
                    namespace = tag['Value']
                    if 'free-tier-' in namespace:
                        env_name = namespace.replace(
                            'free-tier-', '').split('-')[0]
                        is_devtier = True
                    else:
                        break
                if 'kubernetes.io/created-for/pvc/name' in tag['Key']:
                    if 'tiflash-' in tag['Value']:
                        component_name = 'tiflash'
                    if 'tikv-' in tag['Value']:
                        component_name = 'tikv'
                    if 'pd-' in tag['Value']:
                        component_name = 'pd'
                    if 'tidbmonitor-' in tag['Value']:
                        component_name = 'tidbmonitor'
                if 'servicetype' in tag['Key']:
                    # has been labeled
                    break
        if is_devtier:
            ec2client.create_tags(Resources=[volume['VolumeId']], Tags=[{'Key': 'component', 'Value': component_name}, {'Key': 'servicetype', 'Value': 'devtier'}, {
                'Key': 'environment', 'Value': env_name}, {'Key': 'project', 'Value': ''}, {'Key': 'tenant', 'Value': ''}, {'Key': 'cluster', 'Value': ''}, {'Key': 'usedby', 'Value': ''}])


def tag_elb(elbv2, region, cluster_list):
    config = Config(
        retries={
            'max_attempts': 10,
            'mode': 'adaptive'
        }
    )
    elbclient = boto3.client('elbv2', region_name=region, config=config)
    # v2
    for lb in elbv2:
        cluster_name = ''
        is_devtier = False
        is_labeled = False
        if 'Tags' in lb:
            for tags in lb['Tags']:
                for tag in tags['Tags']:
                    if 'elbv2.k8s.aws/cluster' in tag['Key']:
                        cluster_name = tag['Value']
                        if cluster_name in cluster_list:
                            is_devtier = True
                    if 'servicetype' in tag['Key']:
                        # has been labeled
                        is_labeled = True
                        break
                if is_labeled:
                    break
            if is_labeled:
                continue
        if is_devtier:
            elbclient.add_tags(
                ResourceArns=[
                    lb['LoadBalancerArn'],
                ],
                Tags=[
                    {'Key': 'servicetype', 'Value': 'devtier'},
                    {'Key': 'environment',
                        'Value': cluster_name.split('-')[0]},
                    {'Key': 'usedby', 'Value': cluster_name},
                    {'Key': 'component', 'Value': 'tidb'},
                    {'Key': 'project', 'Value': ''},
                    {'Key': 'tenant', 'Value': ''},
                    {'Key': 'cluster', 'Value': ''},
                ]
            )


def tag_eks_vpc_nat_igw(cluster_list):
    for eks_cluster_name in cluster_list:
        region = '-'.join(eks_cluster_name.split('-')[1:-1])
        eks_client = boto3.client('eks', region_name=region)
        ec2_client = boto3.client('ec2', region_name=region)
        ec2_resources_client = boto3.resource('ec2', region_name=region)

        # EKS
        eks_cluster = eks_client.describe_cluster(
            name=eks_cluster_name)['cluster']
        eks_arn = eks_cluster['arn']
        eks_client.tag_resource(
            resourceArn=eks_arn,
            tags={
                'servicetype': 'devtier',
                'component': 'devtier-control-plane',
                'environment': eks_cluster_name.split('-')[0],
                'tenant': 'admin',
                'usedby': eks_cluster_name
            }
        )

        eks_vpc_id = eks_cluster['resourcesVpcConfig']['vpcId']

        # VPC
        vpc = ec2_resources_client.Vpc(eks_vpc_id)
        vpc.create_tags(
            Tags=[
                {'Key': 'servicetype', 'Value': 'devtier'},
                {'Key': 'environment',
                    'Value': eks_cluster_name.split('-')[0]},
                {'Key': 'usedby', 'Value': eks_cluster_name},
                {'Key': 'component', 'Value': 'devtier-control-plane'},
                {'Key': 'project', 'Value': ''},
                {'Key': 'tenant', 'Value': 'admin'},
                {'Key': 'cluster', 'Value': ''},
            ]
        )

        # NAT
        nats = ec2_client.describe_nat_gateways(
            Filters=[
                {
                    'Name': 'vpc-id',
                    'Values': [
                        eks_vpc_id,
                    ]
                },
            ],
        )['NatGateways']
        for nat in nats:
            ec2 = boto3.client('ec2', region_name=region)
            ec2.create_tags(Resources=[nat['NatGatewayId']], Tags=[{'Key': 'component', 'Value': 'devtier-control-plane'}, {'Key': 'servicetype', 'Value': 'devtier'}, {
                'Key': 'environment', 'Value': eks_cluster_name.split('-')[0]}, {'Key': 'tenant', 'Value': 'admin'}, {'Key': 'usedby', 'Value': "%s" % eks_cluster_name}])

        igws = ec2_client.describe_internet_gateways(
            Filters=[
                {
                    'Name': 'attachment.vpc-id',
                    'Values': [
                        eks_vpc_id,
                    ]
                },
            ],
        )['InternetGateways']

        # IGW
        for igw in igws:
            internet_gateway = ec2_resources_client.InternetGateway(
                igw['InternetGatewayId'])
            internet_gateway.create_tags(
                Tags=[
                    {'Key': 'servicetype', 'Value': 'devtier'},
                    {'Key': 'environment',
                        'Value': eks_cluster_name.split('-')[0]},
                    {'Key': 'usedby', 'Value': eks_cluster_name},
                    {'Key': 'component', 'Value': 'devtier-control-plane'},
                    {'Key': 'project', 'Value': ''},
                    {'Key': 'tenant', 'Value': 'admin'},
                    {'Key': 'cluster', 'Value': ''},
                ]
            )


def tag_kms(region):
    kmsclient = boto3.client('kms', region_name=region)

    for key in kmsclient.list_keys()['Keys']:
        tags = kmsclient.list_resource_tags(KeyId=key['KeyId'])['Tags']
        is_dev_tier = False
        for tag in tags:
            if tag['TagKey'] == 'servicetype':
                break
            if tag['TagKey'] == 'sa' and 'free-tier-' in tag['TagValue']:
                is_dev_tier = True
        if is_dev_tier:
            if kmsclient.describe_key(KeyId=key['KeyId'])['KeyMetadata']['KeyState'] != 'PendingDeletion':
                response = kmsclient.tag_resource(
                    KeyId=key['KeyId'],
                    Tags=[
                        {
                            'TagKey': 'servicetype',
                            'TagValue': 'devtier'
                        },
                    ]
                )
                print(response)


def tag_s3(cluster_list):
    for cluster_name in cluster_list:
        s3 = boto3.resource('s3')
        bucket_tagging = s3.BucketTagging('%s-backup' % cluster_name)
        is_labeled = False
        tag_set = [
            {'Key': 'component', 'Value': 'devtier-control-plane'}, {'Key': 'servicetype', 'Value': 'devtier'}, {
                'Key': 'environment', 'Value': cluster_name.split('-')[0]}, {'Key': 'tenant', 'Value': 'admin'}, {'Key': 'usedby', 'Value': "%s" % cluster_name}
        ]
        try:
            for tag in bucket_tagging.tag_set:
                if tag['Key'] == 'servicetype':
                    is_labeled = True
                    break
            if is_labeled:
                continue
            bucket_tagging.put(
                Tagging={
                    'TagSet': tag_set + bucket_tagging.tag_set
                },
            )
        except botocore.exceptions.ClientError:
            bucket_tagging.put(
                Tagging={
                    'TagSet': tag_set
                },
            )


def analysis_resources_list(resources):
    cluster_list = os.getenv('CLUSTER_LIST').split(',')
    if os.getenv('TAG_STATIC') == 'True':
        tag_eks_vpc_nat_igw(cluster_list)
        print("tag eks, vpc, nat, igw")
        tag_s3(cluster_list)
        print("tag s3")
        for resource in resources:
            tag_ec2(
                resource['ec2'], resource['region'], cluster_list)
        print("tag ec2")
    if os.getenv('TAG_DYNAMIC') == 'True':
        for resource in resources:
            tag_ebs(
                resource['ebs'], resource['region'])
            tag_elb(
                resource['elbv2'], resource['region'], cluster_list)
        print("tag ebs, elb")
    if os.getenv('TAG_KMS') == 'True':
        for resource in resources:
            tag_kms(resource['region'])
        print("tag kms")


def crawl_aws_resources(region):
    ec2client = boto3.client('ec2', region_name=region)
    resources_list = {'region': region, 'ec2': [],
                      'ebs': [], 'elbv2': [], 'vpc': [], 'kms': [], 'nat': [], 'igw': [], 'eks': []}
    resources_list['ec2'] = ec2client.describe_instances()
    resources_list['ebs'] = ec2client.describe_volumes()
    resources_list['elbv2'] = get_aws_resources_elb_v2(region)
    return resources_list


class DatetimeEncoder(json.JSONEncoder):
    def default(self, obj):
        try:
            return super().default(obj)
        except TypeError:
            return str(obj)


def crawl_all_region_aws_resources():
    result = []
    # for region in get_available_region_list():
    for region in os.getenv('REGION_LIST').split(','):
        result.append(crawl_aws_resources(region))
    return result


resources = crawl_all_region_aws_resources()
analysis_resources_list(resources)