#!/usr/bin/env python

import os
import json
import subprocess

cluster_id = os.environ['CLUSTER_ID']
project_id = os.environ['PROJECT_ID']
env = os.environ['ENV'] # nightly-ms, staging-ms, prod
current_env = os.environ.copy()

kubectl = '/usr/local/bin/kubectl'
aws = '/usr/local/bin/aws'
gcloud = '/google-cloud-sdk/bin/gcloud'
cred_file = current_env['GOOGLE_APPLICATION_CREDENTIALS']
infra_kubeconfig = "--kubeconfig=/kube/infra-api.kubeconfig"
garden_kubeconfig = "--kubeconfig=/kube/gardener.kubeconfig"

# get cluster
cluster_output = subprocess.check_output([kubectl, infra_kubeconfig, 'get', 'cluster', '-n', env+'-'+project_id, cluster_id, '-ojson'], env=current_env).decode('utf-8')
cluster = json.loads(cluster_output)
provider_name = cluster['metadata']['labels']['cloud.pingcap.com/provider-type']
region = cluster['metadata']['labels']['cloud.pingcap.com/region']
network_name = cluster['spec']['networkName']

# get network
network_output = subprocess.check_output([kubectl, infra_kubeconfig, 'get', 'network', '-n', env+'-'+project_id, network_name, '-ojson'], env=current_env).decode('utf-8')
network = json.loads(network_output)
if provider_name == 'eks' or provider_name == 'gke':
    k8s_name = network['metadata']['annotations']['network.cloud.pingcap.com/dataplane']
    account_name = network['spec']['accountRef']['name']

    # get account info
    account_output = subprocess.check_output([kubectl, infra_kubeconfig, 'get', 'account', account_name, '-ojson'], env=current_env).decode('utf-8')
    account = json.loads(account_output)

    # fetch kubeconfig
    if provider_name == 'eks':
        aws_account = account['spec']['aws']['accountID']
        subprocess.check_output([aws, 'eks', 'update-kubeconfig', '--name', k8s_name, '--region', region, '--profile', aws_account], env=current_env)
    elif provider_name == 'gke':
        gcp_project = account['spec']['gcp']['projectID']
        # if cred_file env is not empty, use it
        if cred_file and os.path.exists(cred_file):
            subprocess.check_output([gcloud, 'auth', 'login', '--cred-file', cred_file], env=current_env)
        subprocess.check_output([gcloud, 'container', 'clusters', 'get-credentials', k8s_name, '--region', region, '--project', gcp_project], env=current_env)
elif provider_name == 'aws' or provider_name == 'gcp': # old arch, it's deploy in the controlplane account
    # get shoot k8s name
    k8s_name = network['status']['fromProvider']['shootName']
    # remove the '-ms' suffix of env and add 'garden-' prefix
    garden_ns = 'garden-' + env.rsplit('-', 1)[0]
    get_secret_process = subprocess.Popen([kubectl, garden_kubeconfig, 'get', 'secret', '-n', garden_ns, k8s_name+'.kubeconfig', '-ojsonpath={.data.kubeconfig}'], stdout=subprocess.PIPE)
    base64_decode_process = subprocess.Popen(['base64', '-d'], stdin=get_secret_process.stdout, stdout=subprocess.PIPE)

    with open('/root/.kube/config', 'wb') as file:
        file.write(base64_decode_process.communicate()[0])
    # subprocess.check_output([kubectl, garden_kubeconfig, 'get', 'secret', '-n', garden_ns, k8s_name+'.kubeconfig', '-ojsonpath={.data.kubeconfig}', '|', 'base64', '-d', '>', '/root/.kube/config'])
else:
    print("unknown provider {provider_name}")

# set kubeconfig namespace
subprocess.check_output([kubectl, 'config', 'set-context', '--current', '--namespace=tidb'+cluster_id])