import subprocess
import json
import logging

import pandas as pd

from dku_nvidia.nim_macro.utils import run
from dku_nvidia.nim_macro.exceptions import NimOperatorError

LOGGER = logging.getLogger(__name__)


# contants
NVIDIA_HELM_REPO_URL = "https://helm.ngc.nvidia.com/nvidia"

GPU_OPERATOR_HELM_NAME = "gpu-operator"
GPU_OPERATOR_HELM_CHART = "nvidia/gpu-operator"

NIM_OPERATOR_HELM_NAME = "nim-operator"
NIM_OPERATOR_HELM_CHART = "nvidia/k8s-nim-operator"


def list(helm, gpu_operator_namespace, nim_operator_namespace):
    """
    Retrieves the status of the GPU and NIM Operator Helms chart deployments.
    """
    LOGGER.info("Listing NIM operators.")
    operator_status = []
    
    # GPU Operator
    cmd = [helm, "list", "-n", gpu_operator_namespace, "-o", "json"]
    err_msg = "Error listing GPU Operator Helm chart: {stderr}"
    r = run(cmd, err_msg, NimOperatorError)
               
    charts = json.loads(r.stdout.decode('utf-8'))
    operator_status += [c for c in charts if c["name"] == GPU_OPERATOR_HELM_NAME]

    # NIM Operator
    cmd = [helm, "list", "-n", nim_operator_namespace, "-o", "json"]
    err_msg = "Error listing NIM Operator Helm chart: {stderr}"
    r = run(cmd, err_msg, NimOperatorError)
               
    charts = json.loads(r.stdout.decode('utf-8'))
    operator_status += [c for c in charts if c["name"] == NIM_OPERATOR_HELM_NAME]

    # Format as HTML table
    df = pd.DataFrame(operator_status)
    html = df.to_html(index=False, justify="center", col_space=150)
    LOGGER.info(f"Found {len(operator_status)} operators.")
    
    return operator_status, html


def add(helm, gpu_operator_namespace, gpu_operator_version, nim_operator_namespace, nim_operator_version):
    """
    Add the GPU and NIM Operators to the cluster.
    This function requires that neither Operator be present on the cluster.
    """
    LOGGER.info("Adding NIM operators.")
    # Ensure that neither the GPU Operator nor NIM Operator are installed on the cluster,
    # as "helm install" is unfortunately not idempotent.
    operators, _ = list(helm, gpu_operator_namespace, nim_operator_namespace)
    if operators:
        LOGGER.warning("Operators already installed, skipping add.")
        return "GPU Operator and/or NIM Operator have previously been installed. Please run the \"Inspect\" action for details"
    
    # Install and upgrade the Nvidia Helm repository
    cmd_repo_add = [helm, "repo", "add", "nvidia", NVIDIA_HELM_REPO_URL]
    cmd_repo_upd = [helm, "repo", "update"]
    err_msg = "Nvidia Helm chart installation failed with error: {stderr}"
    
    LOGGER.info("Adding and updating NVIDIA Helm repository.")
    run(cmd_repo_add, err_msg, NimOperatorError)
    run(cmd_repo_upd, err_msg, NimOperatorError)

    LOGGER.info("Installing GPU and NIM Operators.")
    # Install GPU and NIM Operators
    # GPU operator
    cmd = [
        helm, "install", GPU_OPERATOR_HELM_NAME,
        "--version" , gpu_operator_version,
        "-n", gpu_operator_namespace, "--create-namespace", 
        "--wait",
        GPU_OPERATOR_HELM_CHART
    ]
    err_msg = "GPU Operator installation failed with error: {stderr}."
    run(cmd, err_msg, NimOperatorError)

    # NIM operator
    cmd = [
        helm, "install", NIM_OPERATOR_HELM_NAME,
        "--version" , nim_operator_version,
        "-n", nim_operator_namespace, "--create-namespace", 
        "--wait", 
        NIM_OPERATOR_HELM_CHART
    ]
    err_msg = "NIM Operator installation failed with error: {stderr}."
    run(cmd, err_msg, NimOperatorError)

    # Annotate the NIM Operator in case Horizontal Pod Autoscaling is used
    cmd = [
        "kubectl", "annotate", "svc",
        "k8s-nim-operator-metrics-service", "prometheus.io/scrape=true",
        "-n", nim_operator_namespace
    ]
    err_msg = "Failed to annotate NIM Operator: {stderr}"
    run(cmd, err_msg, NimOperatorError)
    
    # Return the installed operators
    _, html = list(helm, gpu_operator_namespace, nim_operator_namespace)
    html = "The following operators have been deployed:<br>" + html
    LOGGER.info("NIM operators added successfully.")
    
    return html
        

def rm(helm, gpu_operator_namespace, nim_operator_namespace):
    """
    Uninstalls the GPU Operator and/or NIM Operator Helm charts.
    """
    LOGGER.info("Removing NIM operators.")
    operators, _ = list(helm, gpu_operator_namespace, nim_operator_namespace)
    deleted_operators = []
    
    # Retrieved operators
    for op in operators:
        cmd = [helm, "uninstall", "-n", op["namespace"], op["name"]]
        err_msg = "Error deleting " + op["name"] + " Helm chart: {std_err}"
        run(cmd, err_msg, NimOperatorError)
           
        op["status"] = "deleted"
        deleted_operators.append(op)
    
    # Format as HTML table
    df = pd.DataFrame(deleted_operators)
    html = df.to_html(index=False, justify="center", col_space=150)
    html = "The following operators have been removed:<br>" + html
    LOGGER.info(f"Removed {len(deleted_operators)} operators.")
    
    return html