import nuke
import os
import sys


# If you are rendering locally using NNRetime GUI licenses, you should set the following environment variable.
# The same is true in a farm environment using a site license, because that is basically an unlimited amount of floating GUI licenses.
# What it does is forcing the node to directly checkout a GUI license instead of first searching for a batch render license,
# which will make the loading of the plugin a bit faster overall
# os.environ["PIXELMANIA_NNRETIME_FORCE_RENDER_USING_GUI_LICENSES"] = "1"

# If you have a few render licenses and also a few GUI licenses, you probably don't want the farm to steal the GUI licenses for rendering.
# The following environment variable makes sure the node won't check out a GUI license, but rather error if it can't checkout a render license
# os.environ["PIXELMANIA_NNRETIME_DONT_RENDER_USING_GUI_LICENSES"] = "1"

# You should always up the following environment variable, so JIT compiled kernel caches can fit on disk
os.environ["CUDA_CACHE_MAXSIZE"] = "4294967296" #4Gb, max allowed value


nuke.load("NNRetime")


def _extract_gpu_index(args):
    # supports: --gpu (defaults to 0), or --gpu <int>
    if "--gpu" not in args:
        return None
    i = args.index("--gpu")
    # "--gpu" alone => default 0
    if i == len(args) - 1:
        return 0
    try:
        return int(args[i + 1])
    except Exception:
        return 0


gpu = _extract_gpu_index(nuke.rawArgs)
if gpu is not None:
    os.environ["PIXELMANIA_NUKE_GPU_INDEX"] = str(gpu)


if nuke.GUI:
    nuke_prefs = nuke.toNode("preferences")
    selected_gpu = nuke_prefs.knob("selectedGPUDeviceIndex").getValue()
    if selected_gpu > 0:
        nuke.knobDefault("NNRetime.cuda_deviceid", str(int(selected_gpu) - 1))
    else:
        nuke.knobDefault("NNRetime.cuda_deviceid", str(0))
