'''
	The goal of this application is to generate a power diagram that will help system administrators to:
	- document the power supply architecture
	- easily spot potential power overloads (for example if the power consumption exceeds the capacity of a cable)
	
	This application takes its input from a database, currently in the form of an sql dump, but it could easily be adapted to read directly from a mysql database
'''

import sqlite3


import os
import re
import sys
import pygraphviz


def add_capacity_constraints(capacity1, capacity2):
	"""
	combines 2 capacity constraints (max amperes) together
	
	:param float capacity1: max amperes for the first capacity, None if there are no constraints
	:param float capacity2: max amperes for the second capacity, None if there are no constraints
	:return float: max amperes for the combined capacity, None if there are no constraints
	"""
	if capacity1 is None:
		return capacity2
	else:
		if capacity2 is None:
			return capacity1
		else:
			return min(capacity1, capacity2)
			
def mysql_to_sqlite( mysql_sql_code ):
	"""
	converts a mysql-compatible sql code into a sqlite-ompatible sql code
	
	note: the original code was found on internet, then tweaked
	"""
	content = mysql_sql_code

	# unused commands
	COMMAND_RE = re.compile(r'^(SET).*?;\n$', re.I | re.M | re.S)
	content = COMMAND_RE.sub('', content)

	# table constraints
	TCONS_RE = re.compile(r'\)(\s*(CHARSET|DEFAULT|ENGINE)(=.*?)?\s*)+;', re.I | re.M | re.S)
	content = TCONS_RE.sub(');', content)

	# remove comments
	# content = re.sub(r'^-- Tab[.]', 'toto', content, flags=re.I | re.M | re.S)

	# sqlite doesn't like 'unsigned' as in `ip_address_3` tinyint(3) unsigned NOT NULL default '27',
	content = re.sub(r' unsigned ', ' ', content)

	# sqlite doesn't like 'enum' as in `type` enum('normal','light_out_management') NOT NULL default 'normal',,
	content = re.sub(r' enum\([^\)]*\) ', ' varchar(255) ', content)

	# insert multiple values
	# INSERTVALS_RE = re.compile(r'^(INSERT INTO.*?VALUES)\s*\((.*)\*;', re.I | re.M | re.S)
	INSERTVALS_RE = re.compile(r'^(INSERT INTO.*?VALUES)\s*([^;]*);', re.I | re.M | re.S)
	#INSERTVALS_RE = re.compile(r'^(INSERT INTO.*?VALUES)\s*((\[^\)](\)));$', re.I | re.M | re.S)
	INSERTVALS_SPLIT_RE = re.compile(r'\)\s*,\s*\(', re.I | re.M | re.S)

	def insertvals_replacer(match):
		insert, values = match.groups()
		# print("insert=%s"%insert)
		# print("values=%s"%values)
		values = re.sub('^\s*\(' ,'', values)
		values = re.sub('\)\s*$' ,'', values)
		replacement = ''
		for vals in INSERTVALS_SPLIT_RE.split(values):
			#print("vals=%s"%vals)
			replacement = '%s\n%s (%s);' % (replacement, insert, vals)
		return replacement

	content = INSERTVALS_RE.sub(insertvals_replacer, content)
	return content

class Machine(object):
	"""
	represents a device with input and output power plugs. It could represent a power consumer such as a server (in which case, it has no output plug), or a power distrubtion unit (in which case, it has one input plug and more than one output plugs), or even something else...
	"""
	def __init__(self, name, power_config):
		self.name = name
		self.input_plugs = {}
		self.output_plugs = {}
		self.current_capacity_constraint = None # the maximum amperes in this connection 
		self.power_config = power_config
		self.power_consumption = 0.0

	def get_plug(self, plug_name):
		plugs = {'i': self.input_plugs, 'o':self.output_plugs}[plug_name[0]]
		if plug_name not in plugs:
			plugs[plug_name] = Plug(plug_name, self, self.power_config)
		return plugs[plug_name]

	def get_max_amperes(self):
		capacity = None
		if len(self.input_plugs) > 0:
			capacity = self.input_plugs.values()[0].get_max_amperes()
			capacity = add_capacity_constraints (capacity, self.current_capacity_constraint)
		return capacity

	def get_outgoing_connections(self):
		outgoing_connections = []
		for conn in self.power_config.connections:
			if conn.from_plug.machine == self:
				outgoing_connections.append(conn)
		return outgoing_connections

	def get_power_consumption(self):
 		power_consumption = self.power_consumption
 		for conn in self.get_outgoing_connections():
 			power_consumption += conn.get_power_consumption()
			# print("machine %s : power_consumption += %f" % (self.name, conn.get_power_consumption()))
		return power_consumption


class Plug(object):
	"""
	represents a power plug (input or output) of a device
	"""
	def __init__(self, name, machine, power_config):
		self.name = name
		self.machine = machine
		self.current_capacity_constraint = None # the maximum amperes in this connection 
		self.power_config = power_config

	def __str__(self):
		return self.machine.name + '.' + self.name

	def get_incoming_connection(self):
		return self.power_config.get_connection_to(self)

# 	def get_outgoing_connections(self):
# 		return self.power_config.get_connections_from(self)

	def is_input_plug(self):
		return self.name[0] == 'i'

	def set_current_capacity_constraint(self, max_amps):
		self.current_capacity_constraint = max_amps

	def get_max_amperes(self):
		capacity = None

		if self.is_input_plug():
			in_con = self.get_incoming_connection()
			if in_con:
				# apply incoming connection amperes limitation
				capacity = add_capacity_constraints (capacity, in_con.get_max_amperes())
		else:
			# apply the machine containing this plug's amperes limitation
			capacity = add_capacity_constraints (capacity, self.machine.get_max_amperes())
				
		# apply this plug's amperes limitation
		capacity = add_capacity_constraints (capacity, self.current_capacity_constraint)

		return capacity

# 	def get_power_consumption(self):
# 		power_consumption = 0.0
# 		for conn in self.get_outgoing_connections():
# 			power_consumption += conn.get_power_consumption()
# 		power_consumption += self.get_power_consumption()
# 		return self.from_plug.get_power_consumption()


class Connection(object):
	"""
	a power cable connecting an input power plug to an output power plug
	"""

	def __init__(self, from_plug, to_plug):
		self.from_plug = from_plug
		self.to_plug = to_plug
		self.current_capacity_constraint = None # the maximum amperes in this connection 

	def __str__(self):
		return str(self.from_plug) + ' -> ' + str(self.to_plug)
		
	def get_max_amperes(self):
		capacity = self.from_plug.get_max_amperes()
		capacity = add_capacity_constraints (capacity, self.to_plug.current_capacity_constraint)
		if self.current_capacity_constraint is not None:
			capacity = min(capacity, self.current_capacity_constraint)
		return capacity

	def is_redundancy_cable(self):
		return not (self.to_plug.name == 'i' or self.to_plug.name == 'i1')

	def get_power_consumption(self):
		# at the moment, this program doesn't handle redundant power supplies properly:
		# the energy dragged by a power supply depends whether both power supplies are connnected to the same provider or not
		if self.is_redundancy_cable(): 
			return 0.0  # consider secondary power supplies to drag 0 power
		else:
			return self.to_plug.machine.get_power_consumption()
			

class PowerConfig(object):
	"""
	the description of how machines are connected together (in terms of electric power)
	"""

	def __init__(self, simpa_db_sql_file_path):
		self.machines = {}
		self.connections = []
		self._parse_from_simpa_db_sql_file(simpa_db_sql_file_path)
		
	def _parse_from_simpa_db_sql_file(self, simpa_db_sql_file_path):

		def get_table_attr( cur, table, key_name, key_value, attr_name ):
			attr_value = None
			cur.execute("SELECT "+attr_name+" FROM "+table+" WHERE "+key_name+"='"+key_value+"'")
			rows = cur.fetchall()
			if len(rows) > 0:
				attr_value = rows[0][0]
			return attr_value

		def machine_name_toMachine_spec_id( cur, machine_name ):
			machine_spec_id = get_table_attr( cur, 'machines', 'name', machine_name, 'machine_spec_id' )
			return machine_spec_id

		def machine_spec_id_to_power_consumption( cur, machine_spec_id ):
			power_consumption = get_table_attr( cur, 'machine_spec_to_power_consumption', 'machine_spec_id', machine_spec_id, 'power_consumption' )
			return power_consumption

		def get_plug_type_attr(cur, plug_type, attr_name):
			# INSERT INTO `powerplug_type_desc` (`plug_type_id`, `genre`, `max_amps`) VALUES
			# ('iec60309_blue_pne6h_32a_m', 'm', 32.0),
			attr_value = get_table_attr( cur, 'powerplug_type_desc', 'plug_type_id', plug_type, attr_name )
			return attr_value
			
		def read_plug_capacity(cur, plug):
			plug_capacity = None
			machine_spec_id = machine_name_toMachine_spec_id(cur, plug.machine.name)
			if machine_spec_id is not None:
				# INSERT INTO `powerplug_desc` (`machine_spec_id`, `powerplug_id`, `plug_type`) VALUES
				# ('atos_mpdu_2901382', 'i', 'iec60309_blue_pne6h_32a_m'),				
				cur.execute("SELECT plug_type FROM powerplug_desc WHERE machine_spec_id='%s' AND powerplug_id='%s'" % (machine_spec_id, plug.name))
				rows = cur.fetchall()
				if len(rows) > 0:
					plug_type = rows[0][0]
					# print('plug_type : %s' % plug_type)
					
					plug_capacity = get_plug_type_attr(cur, plug_type, 'max_amps')
					#if plug_capacity:
					#	print('plug_capacity : %f A' % plug_capacity)
			return plug_capacity

		sqliteDbPath=':memory:'  # sqlite-specific special name for a file stored in memory. We could use something like '/tmp/simpadb.sqlite' here but this would make parsing really slow (1 minute instead of 1s), unless either :
		# - proper fix : group of INSERT statements are surrounded by BEGIN and COMMIT (see http://stackoverflow.com/questions/4719836/python-and-sqlite3-adding-thousands-of-rows)
		# - the file is stored on a solid state disk
		try:
			os.remove(sqliteDbPath)
		except:
			pass
		con = sqlite3.connect(sqliteDbPath)
		f = open(simpa_db_sql_file_path, 'r')
		sql = f.read() # watch out for built-in `str`
		#print(sql)
		cur = con.cursor()
		#print(mysql_to_sqlite(sql))
		cur.executescript(mysql_to_sqlite(sql))

		cur.execute("SELECT * FROM machine_to_power")

		rows = cur.fetchall()

		for row in rows:
			(to_plug_as_str, from_plug_as_str, powercordid)=row
			if to_plug_as_str != '':
				conn = self._add_connection(from_plug_as_str, to_plug_as_str)
				for plug in (conn.from_plug, conn.to_plug):
					plug_capacity = read_plug_capacity(cur, plug)
					if plug_capacity is not None:
						plug.set_current_capacity_constraint(plug_capacity)

		cur.execute("SELECT * FROM power_output_specs")
		rows = cur.fetchall()

		for row in rows:
			# print row
			# ('ups1_o1', 16.0),
			(to_plug_as_str, max_amps_as_str)=row
			to_plug = self._get_plug(to_plug_as_str)
			to_plug.set_current_capacity_constraint(float(max_amps_as_str))
			
		for machine in self.machines.values():

			machine_name = machine.name

			if re.match('simpatix.._..', machine_name):
				machine_name = '_'.join(machine_name.split('_')[0:-1])

			# print(machine_name)

			machine_spec_id = machine_name_toMachine_spec_id(cur, machine_name)
			if machine_spec_id is not None:
				power_consumption = machine_spec_id_to_power_consumption(cur, machine_spec_id)
				if power_consumption is not None:
					machine.power_consumption = power_consumption
			

	def get_connection_to(self, to_plug):
		for connection in self.connections:
			if connection.to_plug == to_plug:
				return connection
		return None

	def _get_machine(self, machine_name):
		if machine_name not in self.machines:
			self.machines[machine_name] = Machine(machine_name, self)
		return self.machines[machine_name]
        
	def _get_plug(self, plug_as_str):
		elements = plug_as_str.split('_')
		plug_name = elements[-1]
		machine_name = plug_as_str[0:-(len(plug_name)+1)]
		machine = self._get_machine(machine_name)
		return machine.get_plug(plug_name)
        
	def _add_connection(self, from_plug_as_str, to_plug_as_str):
		from_plug = self._get_plug(from_plug_as_str)
		to_plug = self._get_plug(to_plug_as_str)
		conn = Connection(from_plug, to_plug)
		self.connections.append(conn)
		return conn
		
	def __str__(self):
		s = ''
		for c in self.connections:
			s += str(c) + '\n'
		return s


def power_config_to_svg( power_config, svg_file_path ):
	"""
	creates a svg diagram representing the input power configuration
	
	:param PowerConfig power_config: the input power config
	"""
	graph = pygraphviz.AGraph()
	graph.graph_attr['overlap']='false'
	graph.graph_attr['splines']='true'
	

	graph.edge_attr['colorscheme']='rdylgn9'  # 'brbg11'
	graph.node_attr['shape']='box'
	graph.node_attr['height']=0.3 # default 0.5 inches
	graph.node_attr['fontname']='Helvetica' # default : Times-Roman
	graph.edge_attr['fontsize']=10 # default : 14 pt
	graph.edge_attr['len']=1.5 # default : 1.0
	
	def hotness_to_hsv_color(hotness):
		"""
		:param float hotness: temperature of the wire ratio (0.0 : cold -> 1.0 : hot) 
		"""
		clamped_hotness = max(min(hotness, 1.0), 0.0)
		return "%f, 1.0, 0.8" % ((1.0-clamped_hotness)*0.33)

	#graph.add_node('a')
	#graph.add_node('b')
	#graph.add_edge(u'a',u'b',color='blue')
	#graph.add_edge(u'b',u'a',color='blue')
	for con in power_config.connections:
		# print(con.from_plug.machine.name, con.to_plug.machine.name)
		if not con.is_redundancy_cable():  # don't display redundancy cables, as they might overlap and hide the main one
			power_consumption = con.get_power_consumption()
			amperes = power_consumption/220.0
			color='green'
			capacity = con.get_max_amperes()
			if capacity == None:
				label = '?'
			else:
				label = str(capacity) + 'A'
				if amperes/capacity > 1.0:
					color='red'
				elif amperes/capacity > 0.75:
					color='orange'
				else:
					color='green'
			label = "%.1f/%s" % (amperes, label)
			#color='//%d' % int(9.0-amperes/capacity*8)
			
			color = hotness_to_hsv_color(pow(amperes/capacity, 4.0))
			graph.add_edge(con.from_plug.machine.name, con.to_plug.machine.name, color=color, label=label, penwidth=capacity*0.25)
	graph.layout(prog='twopi')
	graph.draw(svg_file_path)