#!/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 grab_card_ids():
	# This method runs nvidia-smi to grab the card ids, then returns a list
    
    if not os.path.isfile("/bin/nvidia-smi"):
    	print("nvidia-smi binary not found!")
    	exit(1)

    cmd = ["/bin/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():
        line = line.rstrip().split(":")[2]
        pcibus_num = int(re.sub('[.:]', '', line).rstrip("0"),16)
        card = "PCI:0:{}:0".format(str(pcibus_num))
    	cards.append(card)
    return cards

def grab_card_boardname():
	if not os.path.isfile("/bin/nvidia-smi"):
	 	print("nvidia-smi binary not found!")
	 	exit(1)

	cmd = ["/bin/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[filename] = filename
			
			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)
	cards = grab_card_ids()
	write_xorg_conf(cards)


''' with open("xorgconftest.j2") as f:
	template =  Template(f.read())

num_of_monitors = 2

monitors = []
screens = []
busid = ["PCI:0:6:0","PCI:0:7:0","PCI:0:8:0","PCI:0:9:0"]

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

#output = template.render(monitors=monitors,devices=busid)
#print(output)

res = itertools.combinations(busid,2)
listres = list(res)

for i in listres:
	output = template.render(monitors=monitors,devices=i,screens=screens)
	filename = "xorg.conf." + str(listres.index(i))
	print(filename)
	f = open(filename,"w")
	f.write(output)
	f.close() '''