#!/usr/bin/python

import ConfigParser
import getopt
import os
import re
import socket
import sys
import time
import traceback

chg = None

class ChargerError(Exception):
	pass

class TcpConnection(object):

	def __init__(self, sock=None):
		if sock is None:
			self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
			self.sock.settimeout(5)
		else:
			self.sock = sock

	def connect(self, host, port):
		self.sock.connect((host, port))

	def send(self, msg):
		totalsent = 0
		while totalsent < len(msg):
			sent = self.sock.send(msg[totalsent:])
			if sent == 0:
				raise ChargerError("TCP socket connection broken.")
			totalsent = totalsent + sent


	def recv(self):
		try:
			return self.sock.recv(2048)
		except socket.timeout:
			return ""


	def recv1(self, chunked=False, max_length=500):
		chunks = []
		bytes_recd = 0
		while True:
			try:
				chunk = self.sock.recv(max_length)
			except socket.timeout:
				break
			chunks.append(chunk)
			bytes_recd = bytes_recd + len(chunk)

			if not chunked:
				break
			#if len(chunk) < max_length:
			#	break
	
#		print(b''.join(chunks))
		return b''.join(chunks)




class BatteryPackStatus(object):

	def __init__(self, bid, voltage=0, current=0, flags=None):
		self.bid = bid
		self.voltage = voltage
		self.current = current
		self.flags = flags

		self.overcharged = False
		self.terminate_charge = False
		self.overtemp = False

		self.charging = False
		self.initialized = False

		self.decode_flags()

	def decode_flags(self):
		hexflags = int(self.flags, 16)
		if hexflags & 0x8000:
			self.overcharged = True
		if hexflags & 0x4000:
			self.terminate_charge = True
		if hexflags & 0x1000:
			self.overtemp = True
		if not hexflags & 0x40:
			self.charging = True
		if hexflags & 0x80:
			self.initialized = True



class BatteryManager(object):
	def __init__(self, host, port):
		self.conn = TcpConnection()
		self.conn.connect(host, port)
		
		# check if we are talking to an expected power supply
		res = self.conn.recv().strip()
		if not "battery" in res:
			raise ChargerError("Unexpected response from battery manager: {0}".format(res))


	def get_battery_status(self):

		voltages = []
		currents = []
		status = []

		pack_status = []
		env_status = {}

		self.conn.send("@99RQ\r\n")

		# The battery manager has limited TCP buffers so long responses can be
		# sent in multiple chunks. Keep receiving until we have enough data. 
		res = ""
		while len(res.splitlines()) < 58:
			res += self.conn.recv()


		pack_pattern = "(?P<state>[\*#])(?P<batnum>\d+):(?P<serial>\d+),(?P<current>\d+),(?P<voltage>\d+),(?P<cap>\d+),(?P<cycles>\d+),(?P<flags>\w+)"
		env_pattern = "Env:[ ]*(?P<temp1>-*\d+)C,[ ]*(?P<temp2>-*\d+)C,[ ](?P<humidity>\d+)%RH"
		pack_re = re.compile(pack_pattern)
		env_re = re.compile(env_pattern)

		for line in res.splitlines():
			#print(line)
			m = re.match(pack_re, line)
			if m:
				d = m.groupdict()
				pack_status.append(
				  BatteryPackStatus(d['batnum'], int(d['voltage']),
				  int(d['current']), d['flags']))
			else:
				m = re.match(env_re, line)
				if m:
					env_status = m.groupdict()


		return pack_status, env_status


	def _send_cmd(self, cmd, resp_key="done"):
		self.conn.send(cmd + "\r\n")
		count = 0
		while True:
			time.sleep(1.0)
			if resp_key in self.conn.recv():
				break
			count += 1
			if count > 10:
				raise ChargerError("Battery is not responding to '{0}' cmd.".format(cmd))


	def disable_pack(self, packnum):
		if packnum == -1:
#			print ("disabling all packs")
			self._send_cmd("@99AF")
		else:
			print("disabling pack {0}".format(packnum))
			assert packnum >=0 and packnum < 55
			self._send_cmd("@{0:02}PF".format(packnum), resp_key='turning')

	def all_charging(self):
		self._send_cmd("@99AC")

	def reset_watchdog(self):
		self._send_cmd("WDRST", resp_key="reset")

class PowerSupply(object):

	PS_STATE_ON = 'ON'
	PS_STATE_OFF = 'OFF'

	def __init__(self, host, port):
		self.conn = TcpConnection()
		self.conn.connect(host, port)
		
		# check if we are talking to an expected power supply
		self.conn.send("SYST:VERS?")
		res = self.conn.recv().strip()
		if res != "1999.0":
			raise ChargerError("Unexpected response from power supply: {0}".format(res))

	def _send_cmd(self, cmd, retries=3, dtype=str):
		res = None
		err = None
		while retries > 0:
			try:
				self.conn.send(cmd)
				res = dtype(self.conn.recv().strip())
			except ValueError as e:
				retries -= 1
				err = e
				continue
			break

		if res is None:
			raise ChargerError("Error issuing command '{0}' to "
			  "power supply:\n{0}".format(cmd, err))

		return res


	def get_set_voltage(self):
		return self._send_cmd("OUTP:VOLT?", dtype=float)

#		self.conn.send("OUTP:VOLT?")
#		return float(self.conn.recv().strip())

	def get_set_current(self):
		return self._send_cmd("OUTP:CURR?", dtype=float)
#		self.conn.send("OUTP:CURR?")
#		return float(self.conn.recv().strip())

	def get_meas_voltage(self):
		return self._send_cmd("MEAS:VOLT?", dtype=float)
#		self.conn.send("MEAS:VOLT?")
#		return float(self.conn.recv().strip())

	def get_meas_current(self):
		return self._send_cmd("MEAS:CURR?", dtype=float)
#		self.conn.send("MEAS:CURR?")
#		return float(self.conn.recv().strip())

	def get_state(self):
		return self._send_cmd("OUTP:STATE?")
#		self.conn.send("OUTP:STATE?")
#		return self.conn.recv().strip()

	def get_mode(self):
		return self._send_cmd("OUTP:MODE?")
#		self.conn.send("OUTP:MODE?")
#		return self.conn.recv().strip()


	def set_state(self, state):

		assert state == self.PS_STATE_ON or state == self.PS_STATE_OFF

		self.conn.send("OUTP:STATE " + state)
		time.sleep(0.5)
		act_state = self.get_state()
		if act_state != state:
			raise ChargerError("Trying to set output state to {0} but power "
			  "supply returned {1}".format(state, act_state))
	
	def set_voltage(self, voltage):

		assert type(voltage) == float
		assert voltage >= 0 and voltage < 40
		
		self.conn.send("OUTP:VOLT {0:.2f}".format(voltage))
		time.sleep(0.5)
		act_voltage = self.get_set_voltage()
		if act_voltage != voltage:
			raise ChargerError("Trying to set output voltage to {0} but power "
			  "supply returned {1}".format(voltage, act_voltage))

	def set_current(self, current):

		assert type(current) == float
		assert current >= 0 and current < 60
		
		self.conn.send("OUTP:CURR {0:.2f}".format(current))
		time.sleep(0.5)
		act_current = self.get_set_current()
		if act_current != current:
			raise ChargerError("Trying to set output current to {0} but power "
			  "supply returned {1}".format(current, act_current))

#		print("setting current to {0}".format(current))


class Charger(object):

	STATE_OFF = 'off'
	STATE_RAMPUP = 'rampup'
	STATE_CONSTCURR = 'const A'
	STATE_CONSTVOLT = 'const V'
	STATE_FULL = 'full'

	PACKSTATE_CHG = 'pchg'
	PACKSTATE_OFF = 'poff'

	CHG_ON = 'chgon'
	CHG_OFF = 'chgoff'

	EXIT_CHARGED = 0
	EXIT_ERROR = 1


	def __init__(self, config, powersupply, batterymanager):
		self.battery_status = None
		self.chg_ok = True
		self.config = config
		self.state = self.STATE_OFF

		self.ps = powersupply
		self.batman = batterymanager

		self.debug_chg_cur = 0

		try:
			self.read_config()
		except (ConfigParser.NoSectionError, ConfigParser.NoOptionError) as e:
			raise ChargerError("\nInvalid configuration file:\n" + str(e))

	def read_config(self):
		self.chg_voltage = \
		  self.config.getfloat('charger', 'voltage')

		self.start_current = \
		  self.config.getfloat('charger', 'start-current')

		self.max_chg_current = \
		  self.config.getfloat('charger', 'current')

		self.max_pack_current = \
		  self.config.getint('battery', 'max-pack-current')

		self.max_pack_current_diff = \
		  self.config.getint('battery', 'max-pack-current-diff')

		self.max_current_step_rel = \
		  self.config.getfloat('charger', 'max-current-step-rel')

		self.max_current_step_abs = \
		  self.config.getfloat('charger', 'max-current-step-abs')

		self.max_temp = \
		  self.config.getint('battery', 'max-temp')

		self.max_humidity = \
		  self.config.getint('battery', 'max-humidity')

	def get_battery_status(self, fully_charged):

		s, e = self.batman.get_battery_status()

		voltages = []
		currents = []
		tca = []
		charging = []
		idle = []
		temp = 0
		humidity = 0

		for pack in s:
			voltages.append(float(pack.voltage) / 1000)
			currents.append(pack.current)
			if pack.overcharged or pack.terminate_charge:
				tca.append(int(pack.bid))
			elif pack.charging and pack.initialized:
				charging.append(int(pack.bid))
			elif int(pack.bid) not in fully_charged:
				idle.append(int(pack.bid))

		fully_charged.update(tca)

		# don't use packs which have been shut off for min/max/avg calculation
		for p in sorted(fully_charged, reverse=True):
			voltages.pop(p)
			currents.pop(p)
			# need to have something in volt/curr lists to prevent math mayhem
			if not voltages:
				voltages = [0]
				currents = [0]
				

		if e:
			temp = max([int(e['temp1']), int(e['temp2'])])
			humidity = int(e['humidity'])
		
		status = {
			'chg_v': self.ps.get_meas_voltage(),
			'chg_a': self.ps.get_meas_current(),
			'packv_avg': sum(voltages) / len(voltages),
			'packv_min': min(voltages),
			'packv_max': max(voltages),
			'packc_avg': sum(currents) / len(currents),
			'packc_min': min(currents),
			'packc_max': max(currents),
			'packs_idle': idle,
			'packs_chg': charging,
			'packs_full': fully_charged,
			'packs_tca': tca,
			'temp': temp,
			'humidity': humidity,
		}
		return status


	def print_status(self, status, print_head=False):
		if print_head:
			print("")
			print("Charger                   Pack voltage        Pack current     Pack status")
			print("                          (avg/min/max)       (avg/min/max)    (idle/chg/full)")

		chg_str = "{0:5.2f}V {1:4.1f}A [{2}]".format(status['chg_v'], status['chg_a'],
		  self.state) #status['chg_state'])
		packv_str = "{0:.2f}/{1:.2f}/{2:.2f}".format(status['packv_avg'],
		  status['packv_min'], status['packv_max'])
		packc_str = "{0}/{1}/{2}".format(status['packc_avg'],
		  status['packc_min'], status['packc_max'])
		packs_str = "{0}/{1}/{2}".format(len(status['packs_idle']),
		  len(status['packs_chg']), len(status['packs_full']))

		print("{0:<23}   {1:<17}   {2:<14}   {3}".format(chg_str, packv_str,
		  packc_str, packs_str))

		
	def stop(self):
		self.chg_ok = False
		self.batman.disable_pack(-1)
		self.ps.set_state(self.ps.PS_STATE_OFF)


	def charge(self):
		sleeptime = 6.0
		self.chg_ok = True
		ps_current = self.start_current
		count = 0
		fully_charged = set()

		print("Starting charge ...")

		while (self.chg_ok):
			status = self.get_battery_status(fully_charged)
			self.print_status(status, count == 0)
			count = (count + 1) % 16

			# test for generic problematic conditions

			# If the pack current goes over the limit set in the config, cut it
			# in half and start ramping up again. This is either due to a ramp
			# up which is too fast and the max current diff is not catching it
			# a configuration issue or packs being shut off for whatever reason.
			# Worst thing is we toggle between ramping up and cutting in half
			# until we get to const volatge mode. 
			if status['packc_max'] > self.max_pack_current:
				print("max pack current over set limit of {0}mA, "
				  "lowering charge current ...".format(self.max_pack_current))
				self.ps.set_current(ps_current / 2)
				self.state = self.STATE_RAMPUP

			# disable full packs
			for pack in status['packs_tca']:
				self.batman.disable_pack(pack)
				if pack not in fully_charged:
					# When we are switching off packs there is the risk that we
					# overpower the remaining packs. However, since packs should
					# switch off when we are already in constant voltage mode
					# this should not happen in a real world scenario. Let's
					# just be paranoid.
					chg_packs = len(status['packs_chg']) - 1
					if chg_packs and \
					  self.ps.get_meas_current() / chg_packs > self.max_pack_current:
						ps_current = self.max_pack_current * chg_packs
						self.ps.set_current(ps_current)

					fully_charged.add(pack)

			if (self.state == self.STATE_OFF):
				# set all packs to charging
				self.batman.all_charging()

				self.ps.set_current(self.start_current)
				self.ps.set_voltage(self.chg_voltage)
				self.ps.set_state(self.ps.PS_STATE_ON)

				self.state = self.STATE_RAMPUP

			elif (self.state == self.STATE_RAMPUP):
				if status['packc_max'] < self.max_pack_current and \
				  (status['packc_max'] - status['packc_min']) < \
				  self.max_pack_current_diff:

					new_curr_rel = ps_current * self.max_current_step_rel
					new_curr_abs = ps_current + self.max_current_step_abs
					new_curr = round(min([new_curr_rel, new_curr_abs]), 1)

					if new_curr > self.max_chg_current:
						new_curr = self.max_chg_current
						self.state = self.STATE_CONSTCURR

					self.ps.set_current(new_curr)
					#print("setting charge current to {0}".format(new_curr))

					# This might happen if we top-up a battery which is already
					# pretty full.
					if self.ps.get_mode() == "CV":
						self.state = self.STATE_CONSTVOLT

			elif (self.state == self.STATE_CONSTCURR):
				# check if we reached constant voltage mode
				if self.ps.get_mode() == "CV":
						self.state = self.STATE_CONSTVOLT

			elif (self.state == self.STATE_CONSTVOLT):
				if len(status['packs_chg']) == 0 and \
				  status['chg_v'] > 0:
					print("All packs fully charged, exiting ...")
					self.stop()
					return self.EXIT_CHARGED

			ps_current = self.ps.get_set_current()
			self.batman.reset_watchdog()
			time.sleep(sleeptime)
			


def usage():
	print("""
Usage: charger [OPTION]...

   -c hostname                IP address or hostname of charger
   -b hostname                IP address or hostname of battery
   -f file                    path to configuration file

""")

def main():
	cfg_file = 'charge.cfg'
	global chg
	pshost = None
	batmanhost = None

	try:
		opts, args = getopt.getopt(sys.argv[1:], "c:b:f:", [])
	except getopt.GetoptError as err:
		# print help information and exit:
		print str(err)  # will print something like "option -a not recognized"
		usage()
		sys.exit(2)


	for o, a in opts:
		if o == "-f":
			cfg_file = a
		elif o == "-b":
			batmanhost = a
		elif o == "-c":
			pshost = a
		else:
			assert False, "unhandled option"

	if pshost is None:
		print("A hostname for the power supply must be supplied.")
		sys.exit(2)

	if batmanhost is None:
		print("A hostname for the battery must be supplied.")
		sys.exit(2)

	config = ConfigParser.ConfigParser()
	res = config.read([cfg_file])
	if not res:
		raise ChargerError("Cannot find configuration file '{0}'.".format(
		  cfg_file))

	try:
		psport = config.getint("charger", "port")
		batmanport = config.getint("battery", "port")
	except (ConfigParser.NoSectionError, ConfigParser.NoOptionError) as e:
		raise ChargerError("\nInvalid configuration file:\n" + str(e))
	
	try:
		ps = PowerSupply(pshost, psport)
		ps.set_current(0.5)
		ps.set_state(ps.PS_STATE_ON)
	except socket.error as e:
		print("Cannot connect to power supply at {0}:{1}".format(pshost, psport))
		print(e)
		sys.exit(1)

	print("Power supply has been enabled, connecting to battery ...")

	retries = 3
	while retries > 0:
		retries -= 1
		try:
			batman = BatteryManager(batmanhost, batmanport)
			break
		except Exception as e:
			print("Battery not responding, giving it some more time to boot ...")
			time.sleep(1)
	
	if retries == 0:
		print("Cannot connect to battery at {0}:{1}".format(batmanhost, batmanport))
		sys.exit(1)

	chg = Charger(config, ps, batman)
	chg.charge()

if __name__ == "__main__":

	try:
		main()
	except KeyboardInterrupt:
		print("\n\nStopping charge due to user interrupt, cleaning up ...")
	except ChargerError as e:
		print(e)
	except Exception as e:
		print("\n\nGot an unexpected error:")
		traceback.print_exc()
		print("\nStopping charge due to error!")
	finally:
		if chg:
			chg.stop()
	

	print("exiting")