####
# detect the leaked gcp disk on dev environments
# it will be added into gcloud function later
###
import json
import subprocess
import os
import argparse
import yaml
from kubernetes import client, config
import base64
import time
import csv
from datetime import datetime
import boto3

# eksClusters = {
#     "dev-base-us-west-2":           "us-west-2",
#     "dev-seed-us-west-2":           "us-west-2",
#     "stability-base-us-west-2":     "us-west-2",
#     "stability-seed-us-west-2":     "us-west-2",
#     "staging-base-us-west-2":       "us-west-2",
#     "staging-seed-us-west-2":       "us-west-2",
#     "dev-seed-us-east-1":           "us-east-1",
#     "staging-seed-us-east-1":       "us-east-1",
#     "staging-seed-ap-northeast-1":  "ap-northeast-1",
#     "staging-seed-ap-northeast-2":  "ap-northeast-2",
#     "staging-seed-ap-southeast-1":  "ap-southeast-1",
# }

eksClusters = {
    "prod-base-us-west-2":           "us-west-2",
    "prod-seed-us-west-2":           "us-west-2",
    "prod-seed-us-east-1":             "us-east-1",
    "prod-seed-ap-northeast-1":  "ap-northeast-1",
    "prod-seed-ap-northeast-2":  "ap-northeast-2",
    "prod-seed-ap-southeast-1":  "ap-southeast-1",
}

outfile='aws-disks.csv'
csvFile = open(outfile, 'w', newline='')
fieldnames = ['id', 'size', "zone", "takenBy"]
writer = csv.DictWriter(csvFile, fieldnames=fieldnames)
writer.writeheader()

def process_all_shoot_clusters():
    for eks in eksClusters:
        ebsclient = boto3.client('ec2', region_name=eksClusters[eks])
        os.system('aws eks update-kubeconfig --name="{}" --region {}'.format(eks, eksClusters[eks]))
        config.kube_config.load_kube_config()
        apiClient = client.CoreV1Api()
        namespaces = apiClient.list_namespace()
        regionalEbs = []
        for ns in namespaces.items:
            os.system('aws eks update-kubeconfig --name="{}" --region {}'.format(eks, eksClusters[eks]))
            config.kube_config.load_kube_config()
            apiClient = client.CoreV1Api()

            if 'shoot--' not in ns.metadata.name:
                continue
            print("process shoot {}".format(ns.metadata.name))
            try:
                shootKubeconfig = apiClient.read_namespaced_secret('kubecfg', ns.metadata.name)
                shootKubeconfig = shootKubeconfig.data["kubeconfig"]
                shootKubeconfig= base64.b64decode(shootKubeconfig)
                shootKubeconfigFile= open(ns.metadata.name , "w")
                yaml.dump(yaml.load(shootKubeconfig), shootKubeconfigFile)
                shootKubeconfigFile.close()
                config.kube_config.load_kube_config(config_file=ns.metadata.name)
                apiClient = client.CoreV1Api()
                pvs = apiClient.list_persistent_volume()
                for pv in pvs.items:
                    if pv.spec.aws_elastic_block_store is not None:
                        print("process eks {} pv {}".format(eks, pv.metadata.name))
                        seps = str.split(pv.spec.aws_elastic_block_store.volume_id, "/")
                        pvcNamespace = pv.spec.claim_ref.namespace
                        if "tidb" in pvcNamespace:
                            tcClient = client.CustomObjectsApi()
                            tcs = tcClient.list_namespaced_custom_object(
                                group="pingcap.com",
                                version="v1alpha1",
                                namespace=pvcNamespace,
                                plural="tidbclusters")
                            if len(tcs['items']) == 0:
                                regionalEbs.append(
                                    [seps[len(seps) -1 ],
                                     pv.spec.capacity['storage'],
                                     eksClusters[eks],
                                     ns.metadata.name + '/' + pvcNamespace]
                                )
                                # writer.writerow(
                                #     {
                                #         'id': seps[len(seps) -1 ],
                                #         'size': pv.spec.capacity['storage'],
                                #         "zone": eksClusters[eks],
                                #         "takenBy": ns.metadata.name + '/' + pvcNamespace,
                                #     }
                                # )
            except Exception as e:
                print("bypass shoot {} because {}".format(ns.metadata.name, e))

        for ebs in regionalEbs:
            try:
                response = ebsclient.describe_volumes(
                    VolumeIds=[
                        ebs[0],
                    ],
                )
                writer.writerow(
                    {
                        'id': ebs[0],
                        'size': ebs[1],
                        "zone": ebs[2],
                        "takenBy": ebs[3]
                    }
                )
            except Exception as e:
                print("bypass volume {} because {}".format(ebs[0], e))

process_all_shoot_clusters()
csvFile.close()