#!/bin/env python
# -*- coding: utf-8 -*-
# need to install requests
# this script is used to suspend/resume all clusters in a project

# USAGE: python suspend_cluster.py [project_name] [suspend/resume]
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) != 3:
    print("USAGE: python suspend_cluster.py [project_name] [suspend/resume]")
    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]
is_suspend = True if sys.argv[2] == 'suspend' else False

# 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']:
        if item['cluster_type'] == 'DEDICATED':
            clusters.append(item['id'])
    
    return clusters

def suspend_cluster(project_id, cluster_id, is_suspend):
    cluster_url = f"https://api.tidbcloud.com/api/v1beta/projects/{project_id}/clusters/{cluster_id}"
    headers = {'Content-type': 'application/json'}
    json_data = {'config':{'paused': is_suspend}}

    if dry_run:
        json_data['dry_run'] = True
        return json_data
    
    response = requests.patch(cluster_url,json=json_data, 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()
    
    return response.json()


# main
project_id = get_projectid(project_name)
clusters = get_all_clusters(project_id)
print(f"Found {len(clusters)} clusters in project {project_name}")
print(clusters)

for cluster_id in clusters:
    print(f"Processing cluster {cluster_id}")
    ret = suspend_cluster(project_id, cluster_id, is_suspend)
    print(ret)