import time
import json
import boto3
import subprocess
import datetime
import pytz

from googleapiclient import discovery
from oauth2client.client import GoogleCredentials
credentials = GoogleCredentials.get_application_default()
containerClient = discovery.build('container', 'v1', credentials=credentials)
computeClient = discovery.build('compute', 'v1', credentials=credentials)

eks = ["dev-seed-us-west-2", "dev-seed-us-east-1"]
gke = ["dev-seed-us-west1-3137aff", "dev-seed-us-west2-9f7b847"]
env = 'dev'
eksAlive = []
gkeAlive = []
tz = pytz.timezone( 'Asia/Shanghai' )
nowTime = datetime.datetime.now(tz=tz)
weekday = datetime.datetime.today().weekday()

def fetch_clusters_json():
    cmd = "KUBECONFIG=infra-apiserver /usr/local/bin/kubectl get clusters -A -o json"
    inputDomain = subprocess.Popen(cmd, shell=True,  stdout=subprocess.PIPE)
    jsonS = json.dumps(inputDomain.communicate()[0].decode('utf-8'))
    return json.loads(jsonS)

def fetch_alive_kubernetesclusters():
    rawData = fetch_clusters_json()
    clusters = json.loads(rawData)
    for cluster in clusters['items']:
        clusterKind = cluster['spec']['cluster']['tidb']['instanceType']
        if 'freetier' in clusterKind:
            continue
        networkName = cluster['spec']['networkName']
        if 'aws' in networkName:
            if 'us-west-2' in networkName:
                eksAlive.append("dev-seed-us-west-2")
            elif 'us-east-1' in networkName:
                eksAlive.append("dev-seed-us-east-1")
        elif 'gcp' in networkName:
            if 'us-west1' in networkName:
                gkeAlive.append("dev-seed-us-west1-3137aff")
            elif 'us-west2' in networkName:
                gkeAlive.append("dev-seed-us-west2-9f7b847")


def wakeup_all_eks():
    for region in ['us-west-2', 'us-east-1']:
        asgClient = boto3.client('autoscaling', region_name=region )
        response = asgClient.describe_auto_scaling_groups()
        print(response)
        asgs = response['AutoScalingGroups']

        for asg in asgs:
            asgName = asg['AutoScalingGroupName']
            if 'seed' not in asgName:
                continue
            if env not in asgName:
                continue
            maxSize = asg['MaxSize']
            if maxSize == 0:
                print("process asg {}".format(asgName) )
                asgClient.set_desired_capacity(
                        AutoScalingGroupName=asgName,
                        DesiredCapacity=1,
                        HonorCooldown=False
                )


def sleep_aws_eks():
    if nowTime.hour >= 9 \
            and nowTime.hour < 18 \
            and weekday < 5 :
        wakeup_all_eks()
        return
    else:
        for kube in eks:
            if kube not in eksAlive:
                region = kube.replace("dev-seed-", "")
                asgClient = boto3.client('autoscaling', region_name=region )
                response = asgClient.describe_auto_scaling_groups()
                print(response)
                asgs = response['AutoScalingGroups']

                for asg in asgs:
                    asgName = asg['AutoScalingGroupName']
                    if 'seed' not in asgName:
                        continue
                    if env not in asgName:
                        continue

                    print("process asg {}".format(asgName) )
                    asgClient.set_desired_capacity(
                        AutoScalingGroupName=asgName,
                        DesiredCapacity=0,
                        HonorCooldown=False
                    )
                delete_eks_shoot_vms(region)
            else:
                region = kube.replace("dev-seed-", "")
                asgClient = boto3.client('autoscaling', region_name=region )
                response = asgClient.describe_auto_scaling_groups()
                print(response)
                asgs = response['AutoScalingGroups']

                for asg in asgs:
                    asgName = asg['AutoScalingGroupName']
                    if 'seed' not in asgName:
                        continue
                    if env not in asgName:
                        continue

                    maxSize = asg['MaxSize']
                    if maxSize == 0:
                        print("process asg {}".format(asgName) )
                        asgClient.set_desired_capacity(
                            AutoScalingGroupName=asgName,
                            DesiredCapacity=1,
                            HonorCooldown=False
                        )

def delete_eks_shoot_vms(region):
    ec2Client = boto3.client('ec2', region_name=region)
    response = ec2Client.describe_instances()
    for reservation in response['Reservations']:
        for instance in reservation['Instances']:
            instanceName = None
            if 'Tags' not in instance:
                continue
            tags = instance['Tags']
            for tag in tags:
                if tag['Key'] == 'Name':
                    instanceName = tag['Value']

            if instanceName is None:
                continue

            print("process ec2 {}".format(instanceName) )

            if str(instanceName).startswith('shoot--') and 'master' in instanceName and env in instanceName:
                response = ec2Client.terminate_instances(
                    InstanceIds=[
                        instance['InstanceId'],
                    ],
                )

def delete_gke_shoot_vms(region):
    zones = []
    if region == "us-west2":
        zones = [ "us-west2-a", "us-west2-b", "us-west2-c"]
    elif region == "us-west1":
        zones = ["us-west1-a", "us-west1-b", "us-west1-c"]

    for zone in zones:
        request = computeClient.instances().list(project="pingcap-gardener", zone=zone)
        while request is not None:
            response = request.execute()
            time.sleep(100)
            if 'items' not in response:
                continue

            for instance in response['items']:
                try:
                    if "shoot" in instance['name'] and 'master' in instance['name'] and env in instance['name']:
                        print("delete vm ", instance['name'])
                        deleteResponse = computeClient.instances().delete(project="pingcap-gardener", zone=zone, instance=instance['name'])
                        deleteResponse.execute()
                except: pass
            request = computeClient.instances().list_next(previous_request=request, previous_response=response)
            time.sleep(10)

def wakeup_all_gke():
    location = 'projects/pingcap-gardener/locations/' \
               'us-west2/clusters/dev-seed-us-west2-9f7b847'
    nodePoolId = location + '/nodePools/gardener-node-pool-fe8bf71'
    request= containerClient.projects().locations().clusters().nodePools().setSize(
        name=nodePoolId,
        body={'nodeCount': 1},
    )
    request.execute()

    location = 'projects/pingcap-gardener/locations/' \
               'us-west1/clusters/dev-seed-us-west1-3137aff'
    nodePoolId = location + '/nodePools/gardener-node-pool-33f537a'


    request= containerClient.projects().locations().clusters().nodePools().setSize(
        name=nodePoolId,
        body={'nodeCount': 1},
    )
    request.execute()

def sleep_gcloud_gke():
    if nowTime.hour >= 9 \
            and nowTime.hour < 18 \
            and weekday < 5 :
        wakeup_all_gke()
        return
    else:
        for kubernetesCluster in gke:
            if kubernetesCluster not in gkeAlive:
                if "us-west2" in kubernetesCluster:
                    location = 'projects/pingcap-gardener/locations/' \
                               'us-west2/clusters/dev-seed-us-west2-9f7b847'
                    nodePoolId = location + '/nodePools/gardener-node-pool-fe8bf71'
                    request= containerClient.projects().locations().clusters().nodePools().setSize(
                        name=nodePoolId,
                        body={'nodeCount': 0},
                    )
                    request.execute()
                    delete_gke_shoot_vms("us-west2" )
                elif "us-west1" in kubernetesCluster:
                    location = 'projects/pingcap-gardener/locations/' \
                               'us-west1/clusters/dev-seed-us-west1-3137aff'
                    nodePoolId = location + '/nodePools/gardener-node-pool-33f537a'


                    request= containerClient.projects().locations().clusters().nodePools().setSize(
                        name=nodePoolId,
                        body={'nodeCount': 0},
                    )
                    request.execute()

                    delete_gke_shoot_vms("us-west1" )
            else:
                if "us-west2" in kubernetesCluster:
                    location = 'projects/pingcap-gardener/locations/' \
                               'us-west2/clusters/dev-seed-us-west2-9f7b847'
                    nodePoolId = location + '/nodePools/gardener-node-pool-fe8bf71'
                    request= containerClient.projects().locations().clusters().nodePools().setSize(
                        name=nodePoolId,
                        body={'nodeCount': 1},
                    )
                    request.execute()
                elif "us-west1" in kubernetesCluster:
                    location = 'projects/pingcap-gardener/locations/' \
                               'us-west1/clusters/dev-seed-us-west1-3137aff'
                    nodePoolId = location + '/nodePools/gardener-node-pool-33f537a'


                    request= containerClient.projects().locations().clusters().nodePools().setSize(
                        name=nodePoolId,
                        body={'nodeCount': 1},
                    )
                    request.execute()

fetch_alive_kubernetesclusters()
sleep_aws_eks()
sleep_gcloud_gke()