"""
Patch dm_control OpenGL initialization to use GPU ID other than the first available.
The intialization function reads from environment variable:

    EGL_DEVICE_ID

Useful for distributed environments to use multiple GPUs evenly.

Source code: https://github.com/deepmind/dm_control/blob/master/dm_control/_render/pyopengl/egl_renderer.py#L47

Issue and workaround: https://github.com/deepmind/dm_control/issues/118
"""
__all__ = ["patch_dm_headless_egl_display", "EGL_ENV_VAR"]


import os
import atexit
from typing import Optional
import dm_control._render.pyopengl.egl_renderer
from dm_control._render.pyopengl import egl_ext as EGL


EGL_ENV_VAR = "EGL_DEVICE_ID"


def _create_initialized_headless_egl_display(device_id: Optional[int] = None):
    """
    Creates an initialized EGL display directly on a device.

    There are three ways:
    - pass an explicit GPU ID as arg to `patch_dm_headless_egl_display`
    - set `EGL_DEVICE_ID` environment variable to an int
    - if `EGL_DEVICE_ID` is undefined, use the first device in CUDA_VISIBLE_DEVICES
    """
    from OpenGL import error

    if device_id is not None:
        device_id = int(device_id)
    elif os.environ.get(EGL_ENV_VAR, None) is not None:
        device_id = int(os.environ[EGL_ENV_VAR])
    elif os.environ.get("CUDA_VISIBLE_DEVICES", None) is not None:
        device_id = int(os.environ["CUDA_VISIBLE_DEVICES"].split(",")[0])
    else:
        device_id = None

    devices = EGL.eglQueryDevicesEXT()
    if device_id is not None:
        try:
            devices = [devices[device_id]]
        except IndexError:
            raise IndexError(
                f"Your specified device ID {device_id} is out of range, you only have {len(devices)} device(s) available."
            )

    # below is code copied from dm_control
    for device in devices:
        display = EGL.eglGetPlatformDisplayEXT(
            EGL.EGL_PLATFORM_DEVICE_EXT, device, None
        )
        if display != EGL.EGL_NO_DISPLAY and EGL.eglGetError() == EGL.EGL_SUCCESS:
            # `eglInitialize` may or may not raise an exception on failure depending
            # on how PyOpenGL is configured. We therefore catch a `GLError` and also
            # manually check the output of `eglGetError()` here.
            try:
                initialized = EGL.eglInitialize(display, None, None)
            except error.GLError:
                pass
            else:
                if initialized == EGL.EGL_TRUE and EGL.eglGetError() == EGL.EGL_SUCCESS:
                    return display
    return EGL.EGL_NO_DISPLAY


def patch_dm_headless_egl_display(device_id: Optional[int] = None):
    EGL_DISPLAY = _create_initialized_headless_egl_display(device_id)

    if EGL_DISPLAY == EGL.EGL_NO_DISPLAY:
        raise RuntimeError("Cannot initialize a headless EGL display.")
    atexit.register(EGL.eglTerminate, EGL_DISPLAY)

    # monkey patch before DM suite.load() call
    dm_control._render.pyopengl.egl_renderer.EGL_DISPLAY = EGL_DISPLAY
