#!/usr/bin/env python3

import argparse
import logging

import yaml
from kubernetes import client, config, dynamic
import warnings

PROVIDER_LABEL = "cloud.pingcap.com/provider-type"
PROVIDER_EKS = "eks"
RECONCILER_LABEL = "cloud.pingcap.com/reconciler"


def get_ns(obj):
    # every object should contain the metadata
    return obj["metadata"].get("namespace", "default")


def get_name(obj):
    # every object should contain the metadata
    return obj["metadata"]["name"]


def get_whitelist_networks(whitelist):
    """
    get_whitelist_networks gets networks in whitelist
    A sample whitelist file will look like following:
    - namespace: ns1
      name: name1
    - namesapce: ns2
      name: name1
      # ... more networks
    """
    if whitelist == None:
        return set()
    with open(whitelist, "r") as wl:
        nets = yaml.load(wl, Loader=yaml.FullLoader)
        return {(net["namespace"], net["name"]) for net in nets}


def gen_net_cluster_map(clusters):
    """
    gen_net_cluster_map generates the network to cluster map
    with the key is the namespaced name of the network and the
    value is the cluster object.
    """
    return {
        (
            get_ns(cluster),
            cluster["spec"]["networkName"],
        ): cluster
        for cluster in clusters
    }


def get_account(net):
    """get_ccount returns the account name of the given net."""
    try:
        return net["spec"]["accountRef"]["name"]
    except KeyError:
        logging.info(
            f"account not found for network({get_ns(net)}/{get_name(net)})"
        )
        return ""


def is_target_account(account, accounts):
    """
    is_target_account checks if the given `account` is a target account. If
    the given list `accounts` is empty, then all accounts will be considered.
    """
    if not accounts:
        return True
    if account in accounts:
        return True
    return False


def get_removable_networks(net_cli, cluster_cli, whitelist, accounts):
    """
    get_empty_networks will return list of network objects that has
    no TiDB clusters running on it.
    """
    clusters = cluster_cli.get(
        label_selector=f"{PROVIDER_LABEL}={PROVIDER_EKS}"
    ).items
    nets = net_cli.get(label_selector=f"{PROVIDER_LABEL}={PROVIDER_EKS}").items
    whitelist = get_whitelist_networks(whitelist)
    net_cluster_dict = gen_net_cluster_map(clusters)
    rm_nets = []

    for net in nets:
        ns = get_ns(net)
        name = get_name(net)
        act = get_account(net)
        if (
            is_target_account(act, accounts)
            and (ns, name) not in whitelist
            and (
                ns,
                name,
            )
            not in net_cluster_dict
        ):
            rm_nets.append(net)

    return rm_nets


def remove_backups(backup_cli, net):
    """remove_backups removes all backups under the same namespace of the net."""
    # get all backups associated with the network
    net_ns = get_ns(net)
    bps = backup_cli.get(namespace=net_ns).items
    for bp in bps:
        bp_name = get_name(bp)
        try:
            backup_cli.delete(
                name=bp_name, namespace=net_ns, async_request=True
            )
            logging.info(f"bp({net_ns}/{bp_name}) is deleted")
        except Exception as e:
            logging.error(
                f"failed to delete the vpcendpoint({net_ns}/{bp_name}):{e}"
            )


def remove_vpcendpoints(vpcendpoint_cli, net, vps):
    """remove_vpcendpoints removes all vpcendpoints associated with the net."""
    net_vpc = net["spec"]["forProvider"]["vpcId"]
    for vp in vps:
        vp_name = get_name(vp)
        vp_vpc = vp["spec"]["forProvider"]["vpcId"]
        if vp_vpc == net_vpc:
            logging.info(f"vpcendpoint({vp_name}) is deleted")
            try:
                if vp_vpc == net_vpc:
                    vpcendpoint_cli.delete(name=vp_name, async_request=True)
                logging.info(f"vpcendpoint({vp_name}) is deleted")
            except Exception as e:
                logging.error(
                    f"failed to delete the vpcendpoint({vp_name}):{e}"
                )


def remove_networkpeers(networkpeer_cli, net):
    """remove_networkpeers removes all network peers associated with the net."""
    net_ns, net_name = get_ns(net), get_name(net)
    nps = networkpeer_cli.get(namespace=net_ns).items
    for np in nps:
        np_name = get_name(np)
        # The vpc peering request will be initialized from the TiDB cloud side
        source = np["spec"]["source"]
        if source == net_name:
            try:
                networkpeer_cli.delete(
                    name=np_name,
                    namespace=net_ns,
                    async_request=True,
                )
                logging.info(f"networkpeer({net_ns}/{np_name}) is deleted")
            except Exception as e:
                logging.error(
                    f"failed to delete the networkpeer({net_ns}/{np_name}):{e}"
                )


def remove_reconciler_label(net_cli, net):
    """remove_reconciler_label removes the reconciler label if exist."""
    if "labels" not in net["metadata"].keys():
        # return immediately if the network does not contain any label.
        return
    old_rcsl = net["metadata"]["labels"].pop(RECONCILER_LABEL, None)
    net_ns, net_name = get_ns(net), get_name(net)
    if not old_rcsl:
        # return immediately if the reconciler label does not exist.
        return
    logging.info(
        f"change reconciler({old_rcsl} -> nightly-eks-provider) for network({net_ns}/{net_name})"
    )
    body = {"metadata": {"labels": {RECONCILER_LABEL: None}}}
    net_cli.patch(
        body,
        namespace=net_ns,
        name=net_name,
        content_type="application/merge-patch+json",
    )


def remove_network(net_cli, net):
    """
    remove_network removes the network using dynamic client.
    see https://github.com/kubernetes-client/python-base/blob/3815cdaca8857931a275460f68593b9f4cb78813/dynamic/client.py#L212
    for the full list of args.
    """
    net_ns, net_name = get_ns(net), get_name(net)
    # to ensure the nightly-eks-provider will take over the network object,
    # remove the reconciler label.
    try:
        remove_reconciler_label(net_cli, net)
    except Exception as e:
        logging.error(
            f"failed to change the reconciler of the network({net_ns}/{net_name}): {e}"
        )
        return

    # delete the network object
    try:
        net_cli.delete(name=net_name, namespace=net_ns, async_request=True)
        logging.info(f"network({net_ns}/{net_name}) is deleted")
    except Exception as e:
        # capture the exception, as we want to remove other network
        logging.error(f"failed to delete the network({net_ns}/{net_name}): {e}")


def main(kubeconfig, whitelist, ignorewarn, accounts, dryrun):
    """main is the main entry of the script."""
    if ignorewarn:
        warnings.simplefilter("ignore")
    logging.getLogger().setLevel(logging.INFO)
    config.load_kube_config(kubeconfig)
    cfg = client.Configuration.get_default_copy()
    cfg.verify_ssl = False
    dyn_cli = dynamic.DynamicClient(client.ApiClient(configuration=cfg))
    net_cli = dyn_cli.resources.get(
        api_version="cloud.pingcap.com/v1", kind="Network"
    )
    backup_cli = dyn_cli.resources.get(
        api_version="cloud.pingcap.com/v1", kind="Backup"
    )
    cluster_cli = dyn_cli.resources.get(
        api_version="cloud.pingcap.com/v1", kind="Cluster"
    )
    networkpeer_cli = dyn_cli.resources.get(
        api_version="cloud.pingcap.com/v1", kind="NetworkPeer"
    )
    vpcendpoint_cli = dyn_cli.resources.get(
        api_version="ec2.aws.crossplane.io/v1alpha1", kind="VPCEndpoint"
    )

    # get the vpc endpoint list
    vpe_lst = vpcendpoint_cli.get().items
    rns = get_removable_networks(net_cli, cluster_cli, whitelist, accounts)
    if dryrun:
      print('Networks to remove:')
      for rn in rns:
        print(get_ns(rn), get_name(rn), get_account(rn), rn['metadata']['creationTimestamp'])
      return
    for rn in rns:
      remove_backups(backup_cli, rn)
      remove_networkpeers(networkpeer_cli, rn)
      remove_vpcendpoints(vpcendpoint_cli, rn, vpe_lst)
      remove_network(net_cli, rn)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="""
        infra_objects_recycler removes the aws cloud
        resources by deleting the backup, vpcpeers, vpcendpoints,
        and network objects in order.
        """
    )
    parser.add_argument(
        "--kubeconfig",
        type=str,
        help="path to the kubeconfig file of the infra apiserver",
        required=True,
    )
    parser.add_argument(
        "--whitelist",
        type=str,
        help="path to the network whitelist, networks on the whitelist won't be deleted",
        required=False,
    )
    parser.add_argument(
        "--ignorewarn",
        type=bool,
        help="if true ignore the system warning",
        required=False,
        default=True,
    )
    parser.add_argument(
        "--accounts",
        nargs="*",
        help="target accounts, networks belonging to these accounts will be deleted",
        required=False,
    )
    parser.add_argument(
        "--dryrun",
        type=bool,
        help="only print networks to delete, do not take action",
        action=argparse.BooleanOptionalAction,
        required=False,
        default=True,
    )
    args = parser.parse_args()
    main(args.kubeconfig, args.whitelist, args.ignorewarn, args.accounts, args.dryrun)