#!/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


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']:
        clusters.append({'id':item['id'], 'name':item['name'], 'type':item['cluster_type'], 'status':item['status']['cluster_status']} )
    
    return clusters


# 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(clusters)