"""
Meant to be run by blender CLI as follows:
blender --blender_argument_1 --blender_argument_2 ... --python blender.py --blender_argument_m ... --blender_argument_last -- python_argument_1 python_argument_2 ...
The input argument to this script is analogous to the standard `--cycles-device` argument.
Allowed inputs: 'CUDA', 'OPTIX', 'HIP', 'ONEAPI', 'METAL', 'CUDA+CPU', 'OPTIX+CPU', 'HIP+CPU', 'ONEAPI+CPU', 'METAL+CPU', 'CPU', 'CPU+CPU'
Inputs 'CPU' and 'CPU+CPU' are equivalent.
"""


from sys import argv
import bpy


GPU_DEVICES = ['CUDA', 'OPTIX', 'HIP', 'ONEAPI', 'METAL']
LOGFILE = 'log.out'  # Use `LOGFILE = ''` to disable logging


def separate_python_arguments(all_blender_args: list[str]) -> list[str]:
    """
    When this script is run by the command
    blender --blender_argument_1 ... --python blender.py --blender_argument_m ... -- python_argument_1 python_argument_2 ...
    this function does the following.
    Input: ['blender', '--blender_argument_1', ..., '--', 'python_argument_1', ...]
    Output: ['python_argument_1', 'python_argument_2', ...]
    """
    first_python_argument_reached = False
    python_arguments = []
    for arg in all_blender_args:
        if first_python_argument_reached:
            python_arguments.append(arg)
        if arg == '--':
            first_python_argument_reached = True
    return python_arguments


def enable_gpus(device_type_combo: str) -> None:
    """
    Sets up Blender to use all GPUs it knows of (if present).
    Previously seen GPUs are taken from user preferences,
    current ones are added by the refresh_devies method.
    Writes log to LOGFILE if LOGFILE is not an empty string.
    Input: Same as the script (see beginning of file)
    """
    use_cpus = device_type_combo.endswith('CPU')
    device_type = device_type_combo.removesuffix('+CPU')
    use_gpus = device_type in GPU_DEVICES
    if device_type not in ['CPU'] + GPU_DEVICES:
        raise RuntimeError('Unsupported device type')

    preferences = bpy.context.preferences
    cycles_preferences = preferences.addons['cycles'].preferences
    cycles_preferences.refresh_devices()
    devices = cycles_preferences.devices
    if not devices:
        raise RuntimeError('Unsupported device type')

    if LOGFILE > '':
        with open(LOGFILE, mode='w', encoding='UTF-8') as f:
            f.write('Device types to use: '+device_type_combo+'\n')
    for device in devices:
        if device.type == 'CPU':
            device.use = use_cpus
        else:
            device.use = use_gpus
            if use_gpus and LOGFILE > '':
                with open(LOGFILE, mode='a', encoding='UTF-8') as f:
                    f.write('Enabled gpu '+device.name+' ('+device.type+') with unique ID '+device.id.split('_')[2]+' (if present)\n')

    cycles_preferences.compute_device_type = device_type
    if use_gpus:
        bpy.context.scene.cycles.device = 'GPU'
    else:
        bpy.context.scene.cycles.device = 'CPU'


# PROGRAM STARTS HERE
python_args = separate_python_arguments(argv)
device_types = python_args[0]
# argument_n = python_args[n] ...

enable_gpus(device_types)