import subprocess
import argparse
import time

def resize_pods(cpu_usage: int, mem_usage: int, sort_num: int) -> None:
    # 获取 node 列表
    node_cmd = f'''kubectl top nodes --no-headers | awk '{{gsub("%",""); print $1,$3,$5}}' | 
              awk '$2 < {cpu_usage} && $3 < {mem_usage} && $1 !~ /regional-control/ {{print $1,$2,$3}}' | 
              sort -k2n -k3n | awk '{{print $1}}' | head -n {sort_num}'''
    node_list = subprocess.check_output(node_cmd, shell=True, text=True).splitlines()
    print(f"node_list: {node_list}")

    # cordon节点
    cordon_cmd = '''kubectl cordon {}'''
    for node in node_list:
        subprocess.run(cordon_cmd.format(node), shell=True, check=True)
        print(f"Cordoned Node {node}")

    # 查找并将该节点上非 daemonset 类型的 Pod 删掉
    pod_cmd = '''kubectl get pods --all-namespaces -o json | 
                jq -r --arg node_name "{node}" '.items[] | select(.spec.nodeName == $node_name) | 
                select(.metadata.ownerReferences[].kind != "DaemonSet") | "\(.metadata.namespace) \(.metadata.name)"' '''
    deletion_cmd = '''kubectl delete pod {} --namespace {} --force --grace-period=0'''
    annotate_cmd = '''kubectl annotate pod {} --namespace {} --overwrite=true cluster-autoscaler.kubernetes.io/safe-to-evict=true'''
    for node in node_list:
        pod_list = subprocess.check_output(pod_cmd.format(node=node), shell=True, text=True).splitlines()
        for pod in pod_list:
            pod_namespace, pod_name = pod.split()
            subprocess.run(deletion_cmd.format(pod_name, pod_namespace), shell=True, check=True)
            print(f"Deleted Pod {pod_name} in Namespace {pod_namespace}")
        
        # 休眠60秒，将所有Pod设置成可安全删除
        time.sleep(120)
        for pod in pod_list:
            pod_namespace, pod_name = pod.split()
            subprocess.run(annotate_cmd.format(pod_name, pod_namespace), shell=True, check=True)
            print(f"Annotated Pod {pod_name} in Namespace {pod_namespace} as safely evictable")

def main():
    # 创建解析器对象
    parser = argparse.ArgumentParser()

    # 添加命令行参数
    parser.add_argument("--cpu-usage", type=int, default=50, help="CPU usage threshold")
    parser.add_argument("--mem-usage", type=int, default=50, help="memory usage threshold")
    parser.add_argument("--sort-num", type=int, default=2, help="number of nodes to return")

    # 解析命令行参数
    args = parser.parse_args()

    # 调用封装好的函数
    resize_pods(args.cpu_usage, args.mem_usage, args.sort_num)

if __name__ == "__main__":
    main()