import subprocess
import os.path
import re
import tempfile
from ai2thor.build import COMMIT_ID
import shlex
import shutil

def pci_records():
    records = []
    command = shlex.split('lspci -vmm')
    output = subprocess.check_output(command).decode()

    for devices in output.strip().split("\n\n"):
        record = {}
        records.append(record)
        for row in devices.split("\n"):
            key, value = row.split("\t")
            record[key.split(':')[0]] = value

    return records

def xorg_bus_id():
    bus_id = None
    for r in pci_records():
        if r.get('Vendor', '') == 'NVIDIA Corporation'\
                and r['Class'] in ['VGA compatible controller', '3D controller']:
            bus_id = 'PCI:' + ':'.join(map(lambda x: str(int(x, 16)), re.split(r'[:\.]', r['Slot'])))
            break

    if bus_id is None:
        raise Exception("no valid nvidia device could be found")

    return bus_id


def has_docker():
    with open(os.devnull, "w") as dn:
        return subprocess.call(['which', 'docker'], stdout=dn) == 0


def bridge_gateway():
    output = subprocess.check_output("docker network inspect -f '{{range .IPAM.Config}}{{.Gateway}}{{end}}' bridge", shell=True).decode('ascii').strip()
    if re.match(r'^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$', output):
        return output
    else:
        raise Exception("Didn't receive a single ip address from network inspect bridge: %s" % output)


def nvidia_version():
    version = None
    version_path = '/proc/driver/nvidia/version'
    if os.path.isfile(version_path):
        with open(version_path) as f:
            for line in f:
                if line.startswith('NVRM version: NVIDIA'):
                    match = re.search(r'Kernel Module\s+([0-9\.]+)\s+', line)
                    if match:
                        version = match.group(1)
                        break
    return version


def generate_dockerfile(tag):

    driver_url = 'http://us.download.nvidia.com/XFree86/Linux-x86_64/{version}/NVIDIA-Linux-x86_64-{version}.run'.format(version=nvidia_version())
    driver_filename = os.path.basename(driver_url)

    dockerfile = """
FROM ai2thor/ai2thor-base:{tag}
RUN wget {driver_url} -P /root/
RUN sh /root/{driver_filename} -s --no-kernel-module
""".format(driver_filename=driver_filename, driver_url=driver_url, tag=tag)

    return dockerfile


def image_exists(image_name):
    output = subprocess.check_output("docker images -q %s" % image_name, shell=True)
    return len(output) > 0


def run(image_name, base_dir, command, environment):
    allowed_keys = {'AI2THOR_PORT', 'AI2THOR_CLIENT_TOKEN', 'AI2THOR_HOST'}

    environment_string = ""
    for k,v in environment.items():
        if k in allowed_keys:
            environment_string += " -e %s=%s " % (k,v)

    relative_command = os.path.relpath(command, base_dir)
    docker_command = os.path.join('/root/.ai2thor', relative_command)
    environment_string += " -e %s=%s -e %s='%s'" % ("AI2THOR_DEVICE_BUSID", xorg_bus_id(), "AI2THOR_COMMAND", docker_command)
    command = "docker run -v {base_dir}:/root/.ai2thor -d --privileged {environment} {image_name} /root/start.sh".format(
        environment=environment_string, image_name=image_name, base_dir=base_dir)
    container_id = subprocess.check_output(command, shell=True).decode('ascii').strip()
    return container_id


def kill_container(container_id):
    subprocess.check_output("docker kill %s" % container_id, shell=True)


def build_image():

    version = nvidia_version()
    image_name = 'ai2thor/ai2thor-nvidia-%s:%s' % (version, COMMIT_ID)

    if image_exists(image_name):
        return image_name

    td = tempfile.mkdtemp()
    with open("%s/Dockerfile" % td, "w") as f:
        f.write(generate_dockerfile(tag))

    subprocess.check_call("docker build --rm -t %s %s" % (image_name, td), shell=True)

    shutil.rmtree(td)

    return image_name
