#!/usr/bin/env python

import os
import re
import subprocess
import sys
import time
import urllib

MAX_X_SERVER_STARTUP_TIME = 120
ATTEMPTS = 3
BYTESTOREAD = 1023
STARTING_RESOLUTION = "3840x2160"


def run(cmd, quiet=True):
    if not quiet:
        print "Running: [{}]".format(cmd)
    use_shell = isinstance(cmd, str)
    return subprocess.check_output(cmd, shell=use_shell)

# Convert PCI specificatiapnons like 00000000:04:00.0 into things like PCI:4:0:0
#  since the latter format is what xorg.conf accepts


def convert_busid(busid_long):
    m = re.match(r"([0-9a-fA-F]+):" +
                 r"([0-9a-fA-F]+):" +
                 r"([0-9a-fA-F]+)\." +
                 r"([0-9a-fA-F]+)",
                 busid_long)
    domain, busid, slot, function = (int(x, 16) for x in m.groups())
    xserver_format = "PCI:{}:{}:{}".format(busid, slot, function)
    return xserver_format


run("sudo apt-get update")
run("sudo apt-get -qqy install pkg-config xterm xserver-xorg-video-dummy" +
    " x11-apps mesa-utils xdotool imagemagick x11-utils x11-xserver-utils" +
    " libav-tools")

# grab the first line which should look something like:
# NVRM version: NVIDIA UNIX x86_64 Kernel Module  387.34  Tue Nov 21 03:09:00 PST 2017
with open("/proc/driver/nvidia/version") as version_file:
    version_str = version_file.next()
driver_version = version_str.split()[7]

# Fetch the same version of the NVIDIA runfile as is on the host
# We should really host this ourselves
# self-caching/Ansible-iszation followup tracked at INFRA-5581
success = False
for attempt in range(ATTEMPTS):
    driver_url = ("https://us.download.nvidia.com/XFree86/Linux-x86_64/" +
                  "{0}/NVIDIA-Linux-x86_64-{0}.run").format(driver_version)
    urllib.urlretrieve(driver_url, "/tmp/NVIDIA.run")
    os.chmod("/tmp/NVIDIA.run", 511)  # octal 777
    with open('/tmp/NVIDIA.run', 'rb') as the_script:
        first_char = the_script.read(1)
        if first_char == b'#':
            success = True
            break
        else:
            more_stuff = the_script.read(BYTESTOREAD)
            print "Attempt number {} didn't download a script".format(
                attempt+1)
            print "----------------"
            print "{}{}".format(first_char, more_stuff)
            print "----------------"

if not success:
    exit(1)

run("sudo /tmp/NVIDIA.run --accept-license --no-questions --no-network" +
    " --ui=none --no-kernel-module")

# Pick out the first available GPU's busid (in the PCI:4:0:0 type format)
bus_ids = run("nvidia-smi --query-gpu=pci.bus_id --format=csv,noheader")
bus_ids = bus_ids.split()
gpu_idx = 0
if 'CUDA_VISIBLE_DEVICES' in os.environ:
    cuda_visible_devices = os.environ['CUDA_VISIBLE_DEVICES'].split(',')
    gpu_idx = int(cuda_visible_devices[0])
bus_id_for_xserver = convert_busid(bus_ids[gpu_idx])

# Create an x config file at /etc/X11/XF86Config
# Without all the /dev/nvidia* devices available
#  nvidia-xconfig will still work for this use case
#  but nvidia-xconfig with the --query-gpu-info
#  flag no longer works.
run('sudo nvidia-xconfig --busid="' + bus_id_for_xserver +
    '" --no-probe-all-gpus -o /etc/X11/XF86Config')

os.environ["DISPLAY"] = ":10"

# Actually start the X server! specify a suitable vt device based on the
# GPU number so there are not conflicts
# Runs on display :10
subprocess.Popen('sudo Xorg -noreset +extension GLX +extension' +
                 ' RANDR +extension RENDER ' +
                 '-sharevts -novtswitch -isolateDevice "' + bus_id_for_xserver +
                 '" -logfile /tmp/xdummy.log -config /etc/X11/XF86Config :10 vt41',
                 shell=True)

# Wait for MAX_X_SERVER_STARTUP_TIME seconds or for X to successfully start
# ("xset -q" gives info on the current x settings which is ignored, but the
# return code indicates if it was able to connect)
x_start_counter = 0
while subprocess.call("xset q &>/dev/null", shell=True) != 0:
    if x_start_counter > MAX_X_SERVER_STARTUP_TIME:
        print "X did not start after {} seconds, aborting...".format(
            MAX_X_SERVER_STARTUP_TIME)
        sys.exit(1)
    x_start_counter += 1
    time.sleep(1)

# Set the starting screen resolution to STARTING_RESOLUTION
# xrandr seems to fail even after xset is returning successfully
# without this sleep.  TODO: understand why and find a better check
time.sleep(1)
run("xrandr --fb {}".format(STARTING_RESOLUTION))

# This disables QT from using the shared memory connection to X. This prevents
# a bunch of errors like
# `X Error: BadShmSeg (invalid shared segment parameter) 128`
os.environ["QT_X11_NO_MITSHM"] = "1"

# call whatever is passed in as arguments to this script
subprocess.check_call(" ".join(sys.argv[1:]), shell=True)
