#!/usr/bin/env python
from jinja2 import Template
import itertools
import subprocess
import datetime
import os
import sys
import time
import socket
from subprocess import call
import re
import json

def check_nvidia_smi():
   try:
       smi = subprocess.check_output(["which","nvidia-smi"])
   except subprocess.CalledProcessError:
       print("nvidia-smi binary not found!")
       exit(1)


def grab_card_ids():
    # This method runs nvidia-smi to grab the card ids, then returns a list
    check_nvidia_smi()
    
    cmd = ["nvidia-smi", "--query-gpu=pci.bus_id","--format=csv,noheader"]
    p = subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    
    cards = []

    for line in p.stdout.readlines():
	stripped_line = line.rstrip().split(":")[2]
        #check for different format of pcibus_id. This happens on baremetals
        # i.e. 00000000:06:00.0 not 00000000:00:06.0
        pcibus_id = re.sub('[.:]', '', stripped_line).rstrip("0")
        if not pcibus_id: # empty string, try the other way
                 stripped_line = line.rstrip().split(":")[1]
                 pcibus_id = re.sub('[.:]', '', stripped_line).rstrip("0")
                 if not pcibus_id:
                        print("Error in grab_card_ids: we can not parse the line {}".format(line))
                        print("Command that generated it is: {}".format(cmd))
                        system.exit(1)
        pcibus_num = int(pcibus_id,16)
        card = "PCI:0:{}:0".format(str(pcibus_num))
        cards.append(card)
    return cards

def grab_card_boardname():
    check_nvidia_smi()

    cmd = ["nvidia-smi", "--query-gpu=name","--format=csv,noheader"]
    cards = []
    p = subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    for line in p.stdout.readlines():
        line = line.rstrip()
        cards.append(line)
    return cards

def write_xorg_conf(cards):
    num_of_cards = len(cards) + 1
    boardname = (grab_card_boardname())[0]
    gpus = []

    for i in range(1, num_of_cards):
        monitors = []
        screens = []
        res = list(itertools.combinations(cards,i))

        for j in range(i):
            monitors.append("Monitor"+str(j))
            screens.append("Screen"+str(j))

        for card in res:
            filename = "/etc/X11/xorg.conf." + str(i) + str(res.index(card))
            
            gpus.append({'filename':filename,'boardname':boardname,'monitors':monitors,'screens':screens,'devices':card})


    print(json.dumps(gpus))

if __name__ == "__main__":

    cards = grab_card_ids()
    write_xorg_conf(cards)