#!/usr/bin/env python

import csv
import os
import json
import subprocess

kubectl = '/usr/local/bin/kubectl'
infra_kubeconfig = "--kubeconfig=/kube/infra-api.kubeconfig"
# infra_kubeconfig = "--kubeconfig=/Users/huaxi/.kube/dev-infra-api.kubeconfig"
result_file = 'eks_with_tidb_clusters.csv'

# list clusters managed by eks provider
clusters_output = subprocess.check_output([kubectl, infra_kubeconfig, 'get', 'cluster', '-A', '-l cloud.pingcap.com/provider-type=eks', '-ojson']).decode('utf-8')
clusters = json.loads(clusters_output)
print(len(clusters['items']))

with open(result_file, mode='w', newline='') as file:
    writer = csv.writer(file, delimiter=',', quotechar='"', quoting=csv.QUOTE_ALL)
    for cluster in clusters['items']:
        try:
            cluster_status = cluster.get('status')
            if cluster_status is None:
                continue
            if cluster_status['phase'] != 'Running' and cluster_status['phase'] != 'Suspend':
                continue
            network_name = cluster['spec']['networkName']
            if network_name.startswith('tidbcloud-system'):
                continue
            namespace = cluster['metadata']['namespace']
            if not namespace.startswith('prod') and not namespace.startswith('staging') and not namespace.startswith('nightly'):
                continue

            # get cluster
            cluster_id = cluster['metadata']['name']
            tenant_id = cluster['metadata']['labels']['tags.tidbcloud.com/tenant']
            project_id = cluster['metadata']['labels']['tags.tidbcloud.com/project']

            pd_replicas = cluster['spec']['cluster']['pd']['replicas']
            tidb_replicas = cluster['spec']['cluster']['tidb']['replicas']
            tikv_replicas = cluster['spec']['cluster']['tikv']['replicas']
            tiflash_replicas = cluster['spec']['cluster']['tiflash']['replicas']
            tiproxy_replicas = cluster['spec']['cluster']['tiproxy']['replicas']
            topology = f'{pd_replicas}-pd/{tidb_replicas}-tidb/{tikv_replicas}-tikv/{tiproxy_replicas}-tiproxy/{tiflash_replicas}-tiflash'

            use_self_managed = cluster['metadata']['annotations'].get('tidbcloud.com/use-self-managed-nodegroup') == 'true'
            tidb_nodes_version = ''
            if use_self_managed:
                asg_output = subprocess.check_output([kubectl, infra_kubeconfig, 'get', 'autoscalinggroup', '-l', f'cluster={cluster_id}', '-ojson']).decode('utf-8')
                asgs = json.loads(asg_output)
                for asg in asgs['items']:
                    if tidb_nodes_version == '':
                        tidb_nodes_version = asg['metadata']['annotations']['tidbcloud.com/self-managed-ng-k8s-version']
                    else:
                        break
            else:
                ng_output = subprocess.check_output([kubectl, infra_kubeconfig, 'get', 'nodegroup', '-l', f'cluster={cluster_id}', '-ojson']).decode('utf-8')
                ngs = json.loads(ng_output)
                for ng in ngs['items']:
                  if tidb_nodes_version == '':
                      tidb_nodes_version = ng['status']['atProvider']['version']
                  else:
                      break

            # get network
            network_output = subprocess.check_output([kubectl, infra_kubeconfig, 'get', 'network', '-n', namespace, network_name, '-ojson']).decode('utf-8')
            network = json.loads(network_output)
            account = network['spec']['accountRef']['name']
            k8s_name = network['metadata']['annotations']['network.cloud.pingcap.com/dataplane']
            project_name = network['metadata']['labels'].get('projectName')

            # get eks
            eks_output = subprocess.check_output([kubectl, infra_kubeconfig, 'get', 'cluster.eks', k8s_name, '-ojson']).decode('utf-8')
            eks = json.loads(eks_output)
            version = eks['status']['atProvider']['version']

            # get admin nodegroup
            label_selector = f'tidbcloud-namespace={project_name}-{project_id},tidbcloud-network={network_name}'
            admin_nodegroup_output = subprocess.check_output([kubectl, infra_kubeconfig, 'get', 'nodegroup', '-l', label_selector, '-ojson']).decode('utf-8')
            admin_nodegroups = json.loads(admin_nodegroup_output)
            admin_nodegroup_version = admin_nodegroups['items'][0]['status']['atProvider']['version']

            # get CDC node groups
            cdc_nodegroups = json.loads(subprocess.check_output([kubectl, infra_kubeconfig, 'get', 'nodegroup', '-l', f'component=ticdc,cluster={cluster_id}', '-ojson']).decode('utf-8'))
            with_cdc = False
            cdc_nodegroup_version = 'N/A'
            if len(cdc_nodegroups['items']) > 0:
                with_cdc = True
                cdc_nodegroup_version = cdc_nodegroups['items'][0]['status']['atProvider']['version']

            # get DM node groups
            dm_nodegroups = json.loads(subprocess.check_output([kubectl, infra_kubeconfig, 'get', 'nodegroup', '-l', f'component=tiflow,cluster={cluster_id}', '-ojson']).decode('utf-8'))
            with_dm = False
            dm_nodegroup_version = 'N/A'
            if len(dm_nodegroups['items']) > 0:
                with_dm = True
                dm_nodegroup_version = dm_nodegroups['items'][0]['status']['atProvider']['version']

            writer.writerow([cluster_id, tenant_id, account, k8s_name, version, admin_nodegroup_version, use_self_managed, tidb_nodes_version, topology, with_cdc, cdc_nodegroup_version, with_dm, dm_nodegroup_version])
        except Exception as e:
            print(e)
            continue