#!/bin/env python
# -*- coding: utf-8 -*-
# need to install requests
# this script is used to get all clusters info

# USAGE: python get_cluster_info.py [project_name] 
import os
import sys

import requests
from requests.auth import HTTPDigestAuth
from tabulate import tabulate


def output(response):
    if response.status_code != 200:
        print('Status:', response.status_code, 'Headers:', response.headers, 'Error Response:', response.text)
    else:
        print(response.json())

if len(sys.argv) != 2:
    print("USAGE: python suspend_cluster.py [project_name]")
    exit()

# env variables
tidb_pub = os.environ.get('TIDB_PUB')
tidb_priv = os.environ.get('TIDB_PRIV')
dry_run = os.environ.get('DRY_RUN') or False

# arguments
project_name = sys.argv[1]

# get projects
def get_projectid(project_name):
    project_url = 'https://api.tidbcloud.com/api/v1beta/projects?page_size=30'
    headers = {'Content-type': 'application/json'}

    # get all project and crate dict for looking-up later(project_name: project_id)
    response = requests.get(project_url, headers=headers, auth=HTTPDigestAuth(tidb_pub,tidb_priv))

    project_dict = {}
    if response.status_code != 200:
        print('Status:', response.status_code, 'Headers:', response.headers, 'Error Response:', response.text)
        exit()
    else:
        for project in response.json()['items']:
            project_dict[project['name']] = project['id']

    if project_name not in project_dict:
        print(f"Project {project_name} not found.")
        exit()
    return project_dict[project_name]

def get_all_clusters(project_id):
    cluster_url = f"https://api.tidbcloud.com/api/v1beta/projects/{project_id}/clusters?page_size=30"
    headers = {'Content-type': 'application/json'}

    # get all clusters in the project
    response = requests.get(cluster_url, headers=headers, auth=HTTPDigestAuth(tidb_pub,tidb_priv))

    if response.status_code != 200:
        print('Status:', response.status_code, 'Headers:', response.headers, 'Error Response:', response.text)
        exit()
    
    clusters = []
    for item in response.json()['items']:
        cluster_info = {'id':item['id'], 'name':item['name'], 'type':item['cluster_type'], 'status':item['status']['cluster_status']} 
        nodemap = extract_nodemap(item['status']['node_map'])
        clusters.append(cluster_info | nodemap)
    
    return clusters

def extract_nodemap(nodemap):
    node_list = {'tidb':"0", 'tikv':"0", 'tiflash':"0"}
    
    # if nodemap has tidb entries, update the count
    if 'tidb' in nodemap:
        tidb_count = len(nodemap['tidb'])
        tidb_spec = "--"
        if tidb_count > 0:
            tidb_spec = nodemap['tidb'][0]['node_size']
        node_list['tidb'] = f"{tidb_count} ({tidb_spec})"
        
    if 'tikv' in nodemap:
        tikv_count = len(nodemap['tikv'])
        tikv_spec = "--"
        if tikv_count > 0:
            tikv_spec = nodemap['tikv'][0]['node_size']
        node_list['tikv'] = f"{tikv_count} ({tikv_spec})"
        
    if 'tiflash' in nodemap:
        tiflash_count = len(nodemap['tiflash'])
        tiflash_spec = "--"
        if tiflash_count > 0:
            tiflash_spec = nodemap['tiflash'][0]['node_size']
        node_list['tiflash'] = f"{tiflash_count} ({tiflash_spec})"
 
    return node_list

# main
project_id = get_projectid(project_name)
clusters = get_all_clusters(project_id)
print(f"Found {len(clusters)} clusters in project {project_name}:{project_id}")
print(tabulate(clusters,headers="keys"))