#!/bin/env python
# -*- coding: utf-8 -*-
# need to install requests
# this script is used to create a cluster in a project for one time use.
# it will create a cluster in AWS, TiDB x 1, TiKV x 3 with 8C32G, 200G Storage

# USAGE: python create_cluster.py [cluster_name] [project_name] [region] [initial_password]
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) != 5:
    print("USAGE: python create_cluster.py [cluster_name] [project_name] [region] [initial_password]")
    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
cluster_name = sys.argv[1]
project_name = sys.argv[2]
region = sys.argv[3]
Initial_password = sys.argv[4]

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

# create cluster
print("now creating cluster in project", project_name)
cluster_url = f"https://api.tidbcloud.com/api/v1beta/projects/{project_dict[project_name]}/clusters"

json_data = {'name':cluster_name, 'cluster_type':'DEDICATED', 'cloud_provider':'AWS', 'region': region, 
             'config': {
                 'root_password': Initial_password,
                 'port': 4000,
                 'components': {
                     'tidb': {'node_size': '8C16G', 'node_quantity':1},
                     'tikv': {'node_size': '8C32G', 'node_quantity':3, 'storage_size_gib': 200},
                     # 'tiflash': {'node_size': '8C64G', 'node_quantity':1, 'storage_size_gib': 200},
                 }
             }}

response = requests.post(cluster_url, json=json_data, headers=headers, auth=HTTPDigestAuth(tidb_pub,tidb_priv))
output(response)