#! /usr/bin/python

# -*- python -*-
# -*- coding: utf-8 -*-

#
# TODO:
# just a pkg class with fields for cpu list, and core-thread list
#

debug = 0

import os

def cpustring_to_list(cpustr):
        """Convert a string of numbers to an integer list.

        Given a string of comma-separated numbers and number ranges,
        return a simple sorted list of the integers it represents.

        This function will throw exceptions for badly-formatted strings.

        Returns a list of integers."""

        fields = cpustr.strip().split(",")
        cpu_list = []
        for field in fields:
                ends = field.split("-")
                if len(ends) > 2:
                        raise "Syntax error"
                if len(ends) == 2:
                        cpu_list += range(int(ends[0]), int(ends[1])+1)
                else:
                        cpu_list += [int(ends[0])]

	# remove any duplicates, can't use list(set(cpu_list))
	nlist = []
	for value in cpu_list:
		if value not in nlist:
			nlist.append(value)

        return nlist


def cmpcpu(a, b):
	"""Compare N and M in 'cpuN' and 'cpuM'"""

	return cmp(int(a.name[3:]), int(b.name[3:]))


def cmpcpulist(a, b):
	"""Compare a list of cpus numbers"""

	return cmp(sum(a), sum(b))


class Cpu:
	def __init__(self, basedir, name):
		self.name = name
		self.dir = "%s/%s" % (basedir, name)
		self.reload()

	def readfile(self, name):
		try:
			f = open("%s/%s" % (self.dir, name))
			value = f.readline().strip()
			f.close()
		except:
			raise
		return value

	def reload_online(self):
		self.online = True
		try:
			online_list = cpustring_to_list(self.readfile("../online"))
			if int(self.name[3:]) not in online_list:
				self.online = False

		except:
			pass

	def reload_thread_siblings_list(self):
		try:
			self.thread_siblings_list = \
				cpustring_to_list(self.readfile("topology/thread_siblings_list"))
		except:
			pass

	def reload_core_siblings_list(self):
		try:
			self.core_siblings_list = \
				cpustring_to_list(self.readfile("topology/core_siblings_list"))

		except:
			pass

	def reload_coreid(self):
		try:
			self.coreid = self.readfile("topology/core_id");
		except:
			pass

	def reload_pkgid(self):
		self.pkgid = None
		try:
			if self.online:
				self.pkgid = self.readfile("topology/physical_package_id");
		except:
			pass

	def reload(self):
		self.reload_online()
		self.reload_core_siblings_list()
		self.reload_thread_siblings_list()
		self.reload_coreid();
		self.reload_pkgid();


#
# class Cpus
#
# self.cpus maps cpu names to class cpu instances
# self.sockets maps a socket/pkg id to a list of class cpu instances
# self.cores maps socket/pkg id to a list of HT sibling lists
#
class Cpus:
	def __init__(self, basedir = "/sys/devices/system/cpu"):
		self.basedir = basedir
		self.cpus = {}
		self.sockets = {}
		self.cores = {}
		self.reload()

	def __getitem__(self, key):
		return self.cpus[key]

	def keys(self):
		return self.cpus.keys()

	def has_key(self, key):
		return self.cpus.has_key(key)

	def reload(self):
		for name in os.listdir(self.basedir):
			if name[:3] != "cpu" or not name[3].isdigit():
				continue

			if name in self.cpus:
				self.cpus[name].reload(self.basedir)
			else:
				c = Cpu(self.basedir, name)
				self.cpus[name] = c
				socket = c.pkgid

				# map socket to cpu
				if socket in self.sockets:
					self.sockets[socket].insert(0, c)
				else:
					self.sockets[socket] = [ c, ]

				# map socket to core-thread list
				if socket in self.cores:
					self.cores[socket].insert(0, c.thread_siblings_list)
				else:
					self.cores[socket] = [ c.thread_siblings_list, ]


		# update number of cpus
		self.nr_cpus = len(self.cpus)

		# sort per-socket cpu list by cpu number
		for socket in self.sockets.keys():
			if self.sockets[socket]:
				self.sockets[socket].sort(cmpcpu)

		# remove duplicates from core-thread list, without using set(),
		# and sort by increasing cpu number
		for socket in self.cores.keys():
			if self.cores[socket]:
				nlist = []
				for value in self.cores[socket]:
					if value not in nlist:
						nlist.append(value)
				self.cores[socket] = nlist
				self.cores[socket].sort(cmpcpulist)


def show_cpus(cpus):

	if debug == 0:
		return;

	print "Number pkgs: %d" % len(cpus.sockets)
	print "Number cpus: %d" % cpus.nr_cpus
	print ""

	socks = cpus.sockets.keys()
	socks.sort()

	for socket in socks:
		print "Socket %s" % socket
		print "  Socket cores-threads: "
		print "  %s" % cpus.cores[socket]
		for c in cpus.sockets[socket]:
			print ""
			print "  %s" % c.name
			print "    %s" % c.dir
			print "    online: %s" % c.online
			print "    pkgid : %s" % c.pkgid
			print "    coreid: %s" % c.coreid
			print "    core   siblings: %s" % c.core_siblings_list
			print "    thread siblings: %s" % c.thread_siblings_list
		print ""
		print ""

	print ""
	print ""


if __name__ == '__main__':
	import sys

	cpus = Cpus()
	show_cpus(cpus)

	#
	# There must be enough cpu cores. If not, return -1
	# indicating there is no desirable task->cpu assignment.
	#
	nc = cpus.nr_cpus/4

	ncores = 0
	for socket in cpus.sockets.keys():
		ncores = ncores + len(cpus.cores[socket])

	if nc == 0:
		if debug:
			print "Warning: too few cpus to choose smbd assignment"
		os._exit(2)

	if ncores < nc:
		if debug:
			print "Warning: too few cores to choose smbd assignment"
		os._exit(2)

	#
	# For each socket, choose a cpu from each core until
	# we have found enough cpus. Exhaust all cores in one
	# socket before moving to the next.
	#
	pin = []

	socks = cpus.sockets.keys()
	socks.sort()

	while len(pin) < nc:
		for s in socks:
			while len(cpus.cores[s]) and len(pin) < nc:
				core = cpus.cores[s].pop()
				cpu = max(core)
				pin.append(cpu)
	pin.sort()

	if debug:
		print "Found assignment: %s" % pin

	# Format the string for taskset(1)
	print ",".join(map(lambda x: str(x), pin))
