import os
import time
import shutil
import subprocess
import requests
import json
import logging

import pandas as pd

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

LOGGER = logging.getLogger(__name__)

DOCKER_SECRET_NAME = "ngc-secret"
API_KEY_SECRET_NAME = "ngc-api-secret"
NIM_YAML_FILENAME = "nim.yaml"
NIM_YAML_DIR = "/tmp/nvidia-{unix_timestamp}"
NIM_PORT=8000

NIM_SERVICE_YAML = """
apiVersion: apps.nvidia.com/v1alpha1
kind: NIMService
metadata:
  name: {service_name}
  namespace: {nim_services_namespace}
spec:
  image:
    repository: {url}
    tag: "{version}"
    pullPolicy: IfNotPresent
    pullSecrets:
      - {pull_secret_name}
  authSecret: {api_secret_name}
  env:
  {env_vars_yaml}
  storage:
    pvc:
      create: true
      storageClass: {storage_class}
      size: {volume_size}Gi
      volumeAccessMode: {volume_access_mode}
  replicas: {replicas}
  {node_selector_yaml}
  resources:
    limits:
      nvidia.com/gpu: {num_gpus}
  expose:
    service:
      type: {service_type}
      port: {nim_port} 
    {ingress_yaml}
  {hpa_yaml}
"""

NODE_SELECTOR_YAML = """
  nodeSelector:
    {node_selector}
"""

NIM_ENV_VARS_YAML = """
  - name: NGC_API_ENDPOINT
    value: "{nim_repository_host}"
  - name: NGC_AUTH_ENDPOINT
    value: "{nim_repository_host}"
  - name: NGC_API_SCHEME
    value: {nim_repository_protocol}
"""

ENV_VARS_YAML = """
  - name: {key}
    value: "{value}"
"""

NGINX_INGRESS_YAML = """
    ingress:
      enabled: true
      spec:
        ingressClassName: nginx
        rules:
          - host: {service_name}.{ingress_host}
            http:
              paths:
              - backend:
                  service:
                    name: {service_name}
                    port:
                      number: {nim_port}
                path: {ingress_path}
                pathType: Prefix
"""

HPA_YAML = """
  metrics:
    enabled: true
    serviceMonitor:
      additionalLabels:
        release: kube-prometheus-stack
  scale:
    enabled: true
    hpa:
      maxReplicas: {max_replicas}
      minReplicas: {min_replicas}
      metrics:
      - type: Object
        object:
          metric:
            name: gpu_cache_usage_perc
          describedObject:
            apiVersion: v1
            kind: Service
            name: {service_name}
          target:
            type: Value
            value: "{gpu_cache_usage_perc}"
"""


def add(
    helm,
    nim_services_namespace,
    nim_image_tag,
    storage_class,
    volume_size,
    volume_access_mode,
    use_autoscaler,
    replicas,
    min_replicas,
    max_replicas,
    gpu_cache_usage_perc,
    num_gpus,
    additional_env_vars,
    node_selector,
    exposition_mode,
    ingress_host,
    ingress_path,
    nim_container_registry_host,
    nim_container_registry_username,
    nim_container_registry_api_key,
    override_nim_repository,
    nim_repository_protocol,
    nim_repository_host,
    nim_repository_api_key):
    """
    Add the NIM Service to the cluster.
    This requires that the GPU Operator and NIM Operator be present on the cluster (without them
    the NIM Services don't actually start, but they can be created without error)
    """        
    LOGGER.info("Adding new NIM service.")
    # Ensure NIM Services namespace is created
    create_namespace_if_not_exists(nim_services_namespace)
    
    # Delete API and Docker pull secrets
    # TODO: this can probably be improved, re-creating the secrets doesn't need to happen each time
    LOGGER.info("Deleting existing k8s secrets for NIM service if they exist.")
    delete_k8s_secrets(nim_services_namespace, secrets=[DOCKER_SECRET_NAME, API_KEY_SECRET_NAME])
    
    LOGGER.info("Creating new k8s secrets for NIM service.")
    # Create API and Docker pull secrets
    cmd = [
        "kubectl", "create", "secret", "docker-registry",
        "-n", nim_services_namespace,
        DOCKER_SECRET_NAME, 
        "--docker-server", nim_container_registry_host,
        "--docker-username", nim_container_registry_username,
        "--docker-password", nim_container_registry_api_key
    ]
    err_msg = "Docker pull kubernetes secret creation failed: {stderr}"
    run(cmd, err_msg, NimServiceError)
    
    cmd = [
        "kubectl", "create", "secret", "generic",
        "-n", nim_services_namespace,
        API_KEY_SECRET_NAME, 
        "--from-literal", f"NGC_API_KEY={nim_repository_api_key}"
    ]
    err_msg = "API Key kubernetes secret creation failed: {stderr}"
    run(cmd, err_msg, NimServiceError)

    LOGGER.info("Generating NIM service YAML.")
    # Generate NIM Service YAML
    ## (a) extract URL, version and service name
    url, version, service_name = process_nim_image_tag(nim_image_tag)

    ## (b) replace YAML
    node_selector_yaml = ""
    if node_selector:
        node_selector_yaml = NODE_SELECTOR_YAML.format(node_selector=node_selector)
    
    env_vars_yaml = ""
    if override_nim_repository:
        env_vars_yaml = NIM_ENV_VARS_YAML.format(
            nim_repository_host=nim_repository_host,
            nim_repository_protocol=nim_repository_protocol
        )
    for var in additional_env_vars:
        env_vars_yaml += ENV_VARS_YAML.format(
            key=var["key"],
            value=var["value"]
        )
    
    ingress_yaml = ""
    if exposition_mode=="nginx_ingress":
        service_type = "ClusterIP"
        ingress_yaml = NGINX_INGRESS_YAML.format(
            service_name=service_name,
            ingress_host=ingress_host,
            ingress_path=ingress_path,
            nim_port=NIM_PORT
        )
    elif exposition_mode=="nodeport":
        service_type = "NodePort"
    else:
        raise NIMServiceInstallError(f"NIM Service installation error: 'exposition_mode' {exposition_mode} not recognized.")

    hpa_yaml = ""
    if use_autoscaler:
        hpa_yaml = HPA_YAML.format(
            service_name=service_name,
            min_replicas=min_replicas,
            max_replicas=max_replicas,
            gpu_cache_usage_perc=gpu_cache_usage_perc
        )
        
    nim_service_yaml = NIM_SERVICE_YAML.format(
        service_name=service_name,
        url=url,
        version=version,
        storage_class=storage_class,
        volume_size=volume_size,
        volume_access_mode=volume_access_mode,
        replicas=replicas,
        num_gpus=num_gpus,
        nim_services_namespace=nim_services_namespace,
        nim_port=NIM_PORT,
        api_secret_name=API_KEY_SECRET_NAME,
        pull_secret_name=DOCKER_SECRET_NAME,
        node_selector_yaml=node_selector_yaml,
        env_vars_yaml=env_vars_yaml,
        service_type=service_type,
        ingress_yaml=ingress_yaml,
        hpa_yaml=hpa_yaml
    )

    ## (c) write yaml to a tmp file
    nim_yaml_path = write_nim_yaml_tmpfile(nim_service_yaml)

    LOGGER.info(f"Applying NIM service YAML from {nim_yaml_path}")
    ## Install NIM Service on cluster
    cmd = ["kubectl", "apply", "-f", nim_yaml_path]
    err_msg = "NIM Service installation failed with error: {stderr}"
    run(cmd, err_msg, NimServiceError)
    LOGGER.info(f"NIM Service {service_name} created.")
    
    ## Annotate the NIM Service if using the autoscaler
    time.sleep(5) # it takes a couple seconds for the service to be created
    if use_autoscaler:
        cmd = [
            "kubectl", "annotate", "svc",
            service_name, "prometheus.io/scrape=true",
            "-n", nim_services_namespace
        ]
        err_msg = "Failed to annotate NIM Service {service_name}: {stderr}"
        LOGGER.info(f"Annotating service {service_name} for autoscaler.")
        run(cmd, err_msg, NimServiceError)
        
    LOGGER.info(f"NIM Service {service_name} installed successfully!")
    return f"NIM Service {service_name} installed successfully!"


def list(helm, nim_services_namespace):
    """
    Retrieves the status of the GPU and NIM Operator Helms chart deployments.
    """
    LOGGER.info(f"Listing NIM services in namespace '{nim_services_namespace}'.")
    nim_services = []
    
    # Fetch NIM Services
    cmd = [
        "kubectl", "get", "nimservices.apps.nvidia.com",
        "-n", nim_services_namespace,
        "-o", "json"
    ]
    err_msg = "NIM Service listing failed with error: {stderr}"
    r = run(cmd, err_msg, NimServiceError)

    LOGGER.debug("Parsing NIM services list from kubectl output.")
    # Unpack response
    r = json.loads(r.stdout.decode('utf-8'))
    
    for service in r["items"]:
        nim_services.append({
            "NAME": service["metadata"]["name"],
            "NAMESPACE": nim_services_namespace,
            "LAST DEPLOYED": service["metadata"]["creationTimestamp"],
            "STATUS": service["status"]["state"],
            "ENDPOINT": service["status"].get("model", {}).get("externalEndpoint", "") # ingress endpoint
        })

    # NodePort IP
    # Retrieve pod IPs if not using an ingress
    for service in nim_services:
        if not service["ENDPOINT"] and service["STATUS"]=="Ready": 
            LOGGER.debug(f"Service {service['NAME']} has no endpoint, describing service to find NodePort.")
            cmd = [
                "kubectl", "describe", "svc", service["NAME"],
                "-n", nim_services_namespace
            ]
            err_msg = "Failed to describe k8s svc with error: {stderr}"
            r = run(cmd, err_msg, NimServiceError)
            
            svc_list = r.stdout.decode('utf-8').split()
            service["ENDPOINT"] = "http://" + svc_list[svc_list.index("Endpoints:")+1]
        else:
            service["ENDPOINT"] = "http://" + service["ENDPOINT"] # hacky

    LOGGER.info(f"Found {len(nim_services)} NIM services.")
    # Format as HTML table
    df = pd.DataFrame(nim_services)
    html = df.to_html(index=False, justify="center", col_space=150)
    
    LOGGER.debug("Returning HTML table of NIM services.")
    return html


def rm(helm, nim_services_namespace, nim_service_name):
    """
    Uninstalls the GPU Operator and/or NIM Operator Helm charts.
    """
    LOGGER.info(f"Removing NIM service '{nim_service_name}' from namespace '{nim_services_namespace}'.")
    ## Delete NIM Service
    cmd = [
        "kubectl", "delete", "nimservices.apps.nvidia.com",
        nim_service_name,
        "-n", nim_services_namespace
    ]
    err_msg = "NIM Service deletion failed with error: {stderr}"
    run(cmd, err_msg, NimServiceError)
       
    LOGGER.info(f"NIM Service {nim_service_name} deleted successfully!")
    return f"NIM Service {nim_service_name} deleted successfully!"


#### --------------------------------- ######
####    NIM Service helper functions   ######
#### --------------------------------- ######

def create_namespace_if_not_exists(namespace):
    try:
        r = subprocess.run(["kubectl", "create", "namespace", namespace], capture_output=True) 
        r.check_returncode()
    except Exception:
        LOGGER.warning(f"NIM Services namespace {namespace} creation failed or already exists: {r.stderr.decode('utf-8')}")

def delete_k8s_secrets(namespace, secrets=[]):
    cmd = [
        "kubectl", "delete", "secret", *secrets,
        "-n", namespace,
    ] 
    
    try:        
        r = subprocess.run(cmd, capture_output=True) 
        r.check_returncode()
    except Exception:
        LOGGER.warning(f"Deletion of kubernetes secrets failed, they may not exist: {r.stderr.decode('utf-8')}")
        
def process_nim_image_tag(nim_image_tag):
    """
    Takes the NIM image tag as input, and returns:
    - container URL
    - tag version
    - kubernetes service name
    
    For reference, an example NIM image tag is: 
        nvcr.io/nim/meta/llama-3.1-8b-base:1.1.2
    """
    url = nim_image_tag.split(":")[0] 
    version = nim_image_tag.split(":")[1]
    
    # Generate the 'service name' from the URL
    # TODO: better sanitization for kubernetes service DNS compliance?
    service_name = url.split("/")[-1].replace(".", "")
    
    LOGGER.debug(f"Processed nim_image_tag. url='{url}', version='{version}', service_name='{service_name}'")
    return url, version, service_name

def write_nim_yaml_tmpfile(nim_service_yaml):
    nim_yaml_dir = NIM_YAML_DIR.format(unix_timestamp=int(time.time()))
    LOGGER.debug(f"Creating temporary directory for NIM YAML: {nim_yaml_dir}")
    os.mkdir(nim_yaml_dir)

    nim_yaml_path = nim_yaml_dir + "/" + NIM_YAML_FILENAME
    with open(nim_yaml_path, 'w+') as f:
        f.write(nim_service_yaml)
    
    LOGGER.debug(f"Wrote NIM YAML to {nim_yaml_path}")
    return nim_yaml_path