# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#           http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# --------------------------------
# Please keep distribution of this code to NVIDIA and partners only
# --------------------------------
# EXAMPLE CODE TO GET NGC API KEY FROM NTS
# This script demonstrates how to get an NGC API key using the NVIDIA GPU Cloud (NGC) Token Service.
# The script uses either nvidia-smi or pynvml library to get the GPU device information 
# This device info is sent to an NVIDIA service for validation to exchange for a token
#
# WARNING - this is only sample code please do not use in production, just as an reference
#
# NOTE - plan to update your client code once PDI is available in NVML and nvidia-smi
#        for now the example will hash a UUID into a pdi-like 64 bit hex string 

import re
import subprocess

import pynvml as nvml
import requests


def get_device_info_nvml():
    try:
        nvml.nvmlInit()
        deviceCount = nvml.nvmlDeviceGetCount()
        deviceInfo = []
        for i in range(deviceCount):
            handle = nvml.nvmlDeviceGetHandleByIndex(i)
            uuid = nvml.nvmlDeviceGetUUID(handle)
            name = nvml.nvmlDeviceGetName(handle)
            brand = nvml.nvmlDeviceGetBrand(handle)
            architecture = nvml.nvmlDeviceGetArchitecture(handle)
            pciDeviceId = nvml.nvmlDeviceGetPciInfo(handle).pciDeviceId
            pciDeviceId_64bit = format(pciDeviceId, 'X')
            # pdi = 'Unknown' # nvml.nvmlDeviceGetPDI(handle) # TODO - PDI is not available in NVML API yet
            fakePdi = int(hash(uuid) % 2**64) # hash the uuid to a 16 bit integer
            pdi = fakePdi # TODO - just until PDI is available in NVML API
            pdi_64bit = format(pdi, 'X')
            # Map brand and architecture to human-readable strings
            # brand_str = get_brand_name(brand)
            # architecture_str = get_architecture_name(architecture)

            deviceInfo.append({
                'uuid': uuid,
                'pdi': f"0x{pdi_64bit}",
                'name': name,
                'brand': brand,
                'architecture': architecture,
                'pci_device_id': f"0x{pciDeviceId_64bit}",
            })
        nvml.nvmlShutdown()
        print("NVML device info:")
        print(deviceInfo)
        return deviceInfo
    except nvml.NVMLError as e:
        print(f"NVML Error: {e}")
        nvml.nvmlShutdown()
        return []

def get_device_info_smi():
    try:
        output = subprocess.check_output(['nvidia-smi', '-q'], text=True)
        deviceInfo = []
        uuid_pattern = re.compile(r"GPU UUID\s*:\s*(.+)")
        pdi_pattern = re.compile(r"PDI\s*:\s*(.+)")
        name_pattern = re.compile(r"Product Name\s*:\s*(.+)")
        brand_pattern = re.compile(r"Product Brand\s*:\s*(.+)")
        arch_pattern = re.compile(r"Product Architecture\s*:\s*(.+)")
        pci_pattern = re.compile(r"Device Id\s*:\s*(.+)")

        matches = {
            "uuid": uuid_pattern.search(output),
            "pdi": pdi_pattern.search(output),
            "name": name_pattern.search(output),
            "brand": brand_pattern.search(output),
            "architecture": arch_pattern.search(output),
            "pci_device_id": pci_pattern.search(output),
        }

        # TODO - just until PDI is available in NVML API
        if not matches["pdi"]:
            fakePdi = int(hash(matches["uuid"].group(1).strip()) % 2**64) # hash the uuid to a 16 bit integer for testing
            pdi_64bit = format(fakePdi, 'X')
            matches["pdi"] = re.match(r"(.+)", f"0x{pdi_64bit}")

        deviceInfo.append({k: (m.group(1).strip() if m else "Unknown") for k, m in matches.items()})
        print("SMI device info:")
        print(deviceInfo)
        return deviceInfo
    except subprocess.CalledProcessError as e:
        print(f"nvidia-smi error: {e}")
        return []

def validate_device_info(deviceInfo):
    matchInName = ['RTX', 'OTHER DEVICE NAME TEST']
    for device in deviceInfo:
        if any(keyword in device.get('name', '') for keyword in matchInName):
            print(f"Valid device name {device['name']} found")
            return True
    matchInBrandString = ['GeForce','5']
    for device in deviceInfo:
        if any(keyword in str(device.get('brand', '')) for keyword in matchInBrandString):
            print(f"Valid device brand {device['brand']} found")
            return True
    print("No valid device found in the list")
    return False

def get_ngc_key_from_device_info(deviceInfo):
    # Optional - check env for vars
    # if NGC_API_KEY or NGC_CLI_API_KEY decide if the client should use that instead
    # the key will still need to be setup properly with Catalog Access, and the owner will need an NVAIE sub  # noqa: E501
    # or be a member of the NVIDIA developer program (https://developer.nvidia.com/developer-program)
    # apiKey = os.environ.get('NGC_API_KEY')
    # if apiKey:
    #     print("Using the API key from the environment variable instead of calling API")  # noqa: E501
    #     return apiKey

    ngcKeyServiceUrl = 'https://nts.ngc.nvidia.com/v1/token'
    payload = {
        "client_id": "comfyui",
        "pdi": deviceInfo[0].get('pdi', 'Unknown') if deviceInfo else 'Unknown',
        "access_policy_name": "nim-dev",
        "device": deviceInfo[0] if deviceInfo else {}
    }
    print(f"debug Payload: {payload}")
    try:
        response = requests.post(ngcKeyServiceUrl, headers={'Accept': 'application/json'}, json=payload)  # noqa: E501
        print(f"Response: {response}")
        response.raise_for_status()
        keyData = response.json()
        print(f"Key data: {keyData}")
        return keyData.get('access_token')
    except requests.RequestException as e:
        print(f"Error fetching API key: {e}")
        return None

def get_ngc_key():
    # two different examples of getting device info 
    # note: branch/architecture, nvidia-smi gives strings, NVML gives ints
    deviceInfoNvml = get_device_info_nvml()
    deviceInfoSmi = get_device_info_smi()

    # client check option - TBD if needed by client development team
    # if not validate_device_info(deviceInfoNvml):
    #     print("Device info validation failed - exiting")
    #     exit(1)

    key = get_ngc_key_from_device_info(deviceInfoNvml)
    if key:
        print(f"Key: {key}")
        return key
    else:
        raise Exception("Error getting NGC API key. Please follow the instructions in the README to set up your NGC API key.")  # noqa: E501


if __name__ == "__main__":
    # deviceInfoNvml = get_device_info_nvml()
    # deviceInfoSmi = get_device_info_smi()
    # assert deviceInfoNvml[0]["pdi"] == deviceInfoSmi[0]["pdi"]
    # key = get_ngc_key_from_device_info(deviceInfoSmi)
    get_ngc_key()