import os
import sys
import json
import time
import requests
import os
import googleapiclient
from googleapiclient import discovery
from oauth2client.client import GoogleCredentials
from kubernetes import client, config
from prometheus_client import Gauge, start_http_server

gcp_pd = Gauge('gcp_pd_total_size', 'gcp pd total size')
gcp_ssd = Gauge('gcp_ssd_total_size', 'gcp ssd total size')
gcp_awailable_pd = Gauge('gcp_pd_avail_size', 'gcp pd awailable size')
gcp_awailable_ssd = Gauge('gcp_ssd_avail_size', 'gcp ssd awailable size')
gcp_total_cpus = Gauge('gcp_vms_total_cpus', 'gcp cpus total size')
gcp_total_memorys = Gauge('gcp_vms_total_memorys', 'gcp memory total size')
gcp_tidb_total_cpus = Gauge('gcp_tidb_total_cpus', 'gcp cpus total size')
gcp_tidb_total_memorys = Gauge('gcp_tidb_total_memorys', 'gcp memory total size')
gcp_control_plane_cpus = Gauge('gcp_control_plane_cpus', 'gcp control plane cpu total size')
gcp_control_plane_memorys = Gauge('gcp_control_plane_memorys', 'gcp contral plan memory total size')

gcp_others_cpus = Gauge('gcp_others_cpus', 'gcp others cpu total size')
gcp_others_memorys = Gauge('gcp_others_memorys', 'gcp others[ memory total size')

gcp_shoot_idle_cpus = Gauge('gcp_shoot_idle_cpus', 'gcp shoot idle cpu total size')
gcp_shoot_idle_memorys = Gauge('gcp_shoot_idle_memorys', 'gcp shoot idle memory total size')


projectName = "pingcap-gardener"

credentials = GoogleCredentials.get_application_default()
containerClient = discovery.build('container', 'v1', credentials=credentials)
computeClient = discovery.build('compute', 'v1', credentials=credentials)
allShoots = {}
gcpMachineTypes = {}

def compute_disks( ):
    gcp_pd_size = 0
    gcp_ssd_size = 0
    gcp_pd_available_size = 0
    gcp_ssd_available_size = 0

    request = computeClient.disks().aggregatedList(project = projectName)
    while request is not None:
        response = request.execute()
        if 'items' not in response:
            continue

        zonedInstances = response['items']
        for zones in zonedInstances.items():
            zoneName = zones[0]
            zoneName = str(zoneName).replace('zones/', '')
            disksList = zones[1]
            if 'disks' not in disksList:
                continue
            for disk in disksList['disks']:
                diskName = disk['name']
                size = disk['sizeGb']
                if 'pd-standard' in disk['type']:
                    gcp_pd_size  = gcp_pd_size + int(size)
                    if 'users' not in disk or len( disk['users'] ) == 0:
                        gcp_pd_available_size = gcp_pd_available_size + int(size)
                else:
                    gcp_ssd_size  = gcp_ssd_size + int(size)
                    if 'users' not in disk or len( disk['users'] ) == 0:
                        gcp_ssd_available_size = gcp_ssd_available_size + int(size)

        request = computeClient.disks().aggregatedList_next(previous_request=request, previous_response=response)

    gcp_pd.set(gcp_pd_size)
    gcp_ssd.set(gcp_ssd_size)
    gcp_awailable_pd.set(gcp_pd_available_size)
    gcp_awailable_ssd.set(gcp_ssd_available_size)

def compute_instances():
    gcp_total_cpu_size = 0
    gcp_total_memory_size = 0

    gcp_tidb_total_cpu_size = 0
    gcp_tidb_total_memory_size = 0

    gcp_control_plane_cpu_size = 0
    gcp_control_plane_memory_size = 0

    gcp_others_cpu_size = 0
    gcp_others_memory_size = 0

    gcp_shoot_idle_cpu_size = 0
    gcp_shoot_idle_memory_size = 0
    request = computeClient.instances().aggregatedList(project = projectName)
    while request is not None:
        response = request.execute()
        if 'items' not in response:
            continue

        zonedInstances = response['items']
        for zones in zonedInstances.items():
            zoneName = zones[0]
            zoneName = str(zoneName).replace('zones/', '')

            instancesList = zones[1]
            if 'instances' not in instancesList:
                continue
            for instance in instancesList['instances']:
                # print(instance)
                instanceName = instance['name']
                machineType = instance['machineType']
                machineTypeSeps = machineType.split('/')
                machineType = machineTypeSeps[ len(machineTypeSeps) - 1 ]
                machineZone = machineTypeSeps[ len(machineTypeSeps) - 3 ]
                if machineType not in gcpMachineTypes:
                    machineTypeRequest = computeClient.machineTypes().get(project=projectName,
                                                                          zone=machineZone,
                                                     machineType=machineType)
                    machineTypeResponse = machineTypeRequest.execute()
                    gcpMachineTypes[machineType] = machineTypeResponse

                machineTypeDescription = gcpMachineTypes[machineType]

                cpus = machineTypeDescription['guestCpus']
                memorys = machineTypeDescription['memoryMb']

                gcp_total_cpu_size = gcp_total_cpu_size + cpus
                gcp_total_memory_size = gcp_total_memory_size + memorys / 1024

                if 'seed' in instanceName:
                    gcp_control_plane_cpu_size = gcp_control_plane_cpu_size + cpus
                    gcp_control_plane_memory_size = gcp_control_plane_memory_size + memorys / 1024
                elif 'master' in instanceName and 'shoot--' in instanceName:
                    gcp_shoot_idle_cpu_size = gcp_shoot_idle_cpu_size + cpus
                    gcp_shoot_idle_memory_size = gcp_shoot_idle_memory_size + memorys / 1024
                elif 'shoot--' in instanceName:
                    gcp_tidb_total_cpu_size = gcp_tidb_total_cpu_size + cpus
                    gcp_tidb_total_memory_size = gcp_tidb_total_memory_size + memorys / 1024
                else:
                    gcp_others_cpu_size = gcp_others_cpu_size + cpus
                    gcp_others_memory_size = gcp_others_memory_size + memorys / 1024

        request = computeClient.instances().aggregatedList_next(previous_request=request, previous_response=response)

    gcp_total_cpus.set(gcp_total_cpu_size)
    gcp_total_memorys.set(gcp_total_memory_size)

    gcp_tidb_total_cpus.set(gcp_tidb_total_cpu_size)
    gcp_tidb_total_memorys.set(gcp_tidb_total_memory_size)

    gcp_control_plane_cpus.set(gcp_control_plane_cpu_size)
    gcp_control_plane_memorys.set(gcp_control_plane_memory_size)

    gcp_others_cpus.set(gcp_others_cpu_size)
    gcp_others_memorys.set(gcp_others_memory_size)

    gcp_shoot_idle_cpus.set(gcp_shoot_idle_cpu_size)
    gcp_shoot_idle_memorys.set(gcp_shoot_idle_memory_size)

if __name__ == "__main__":
    start_http_server(8000)
    while True:
        compute_disks()
        compute_instances()
        time.sleep(20)