import dataiku
import os
import subprocess
import requests
import shutil
import logging

LOGGER = logging.getLogger(__name__)

def configure_kubeconfig(cluster_name):
    """
    Sets environment variable for kubeconfig of a cluster.
    """
    client = dataiku.api_client()
    cluster = client.get_cluster(cluster_name)
    cluster_settings = cluster.get_settings().settings
    
    # Don't set if cluster is unmanaged (uses defualt kubeconfig location)
    if not cluster_settings["type"] == "manual":
        LOGGER.info(f"Configuring KUBECONFIG for managed cluster '{cluster_name}'.")
        kubeconfig_path = cluster_settings["data"]["kube_config_path"]
        os.environ["KUBECONFIG"] = kubeconfig_path
    else:
        LOGGER.info(f"Cluster '{cluster_name}' is unmanaged, using default kubeconfig.")
    
    
def get_helm_cmd():
    """
    Retrieve path to Helm command in PATH, 
    or install a local version of Helm if not installed.
    """
    try:
        helm_cmd = subprocess.check_output(["which", "helm"]).strip().decode("utf8")
        LOGGER.info(f"Found helm on the machine at: {helm_cmd}")
    except Exception:
        local_helm_folder = os.path.join(os.environ["DIP_HOME"], "tmp", "local_helm")
        helm_cmd = os.path.join(local_helm_folder+"/linux-amd64", "helm")
        LOGGER.info(f"Helm not found in PATH, will use helm from {local_helm_folder}")

        if not os.path.exists(local_helm_folder):
            os.makedirs(local_helm_folder)

        if not os.path.exists(helm_cmd):
            # Retrieve latest version
            r = requests.get(
                "https://get.helm.sh/helm-latest-version",
                stream=True,
                headers={"User-Agent": "DSS Nvidia Plugin"}
            )
            helm_latest_version = r.content.decode("utf-8").strip()

            # Download helm locally
            LOGGER.info(f"Downloading helm version {helm_latest_version}")
            r = requests.get(
                f"https://get.helm.sh/helm-{helm_latest_version}-linux-amd64.tar.gz",
                stream=True,
                headers={"User-Agent": "DSS Nvidia Plugin"}
            )
            local_helm_archive = os.path.join(local_helm_folder, "helm.tar.gz")
            with open(local_helm_archive, "wb") as f:
                shutil.copyfileobj(r.raw, f)
            subprocess.check_call(["tar", "-xzf", local_helm_archive], cwd=local_helm_folder)
            
    return helm_cmd


def run(cmd, err_msg="", exception=Exception):
    LOGGER.debug(f"Running command: {' '.join(cmd)}")
    try:
        r = subprocess.run(cmd, capture_output=True) 
        r.check_returncode()
    except subprocess.CalledProcessError as err:
        LOGGER.error(f"Command failed. Stderr: {r.stderr.decode('utf-8')}")
        raise exception(err_msg.format(stderr=r.stderr.decode('utf-8'))) from err
        
    return r