#!/usr/bin/python3

import time
import serial
import sys
import numpy as np
from contextlib import contextmanager

class HendiException(Exception):
	pass

class HendiCtrl:
	def __init__(self, port, speed, debug=False):
		self.debug = debug
		self.prompt = b':'
		self.sw_id = None
		self.sw_ver = None

		# Deferred-open pyserial idiom: the port isn't actually opened until
		# connect() is called, so a missing/unplugged device doesn't crash
		# the whole server at construction time (see components/actor/
		# heater_hendi.py's HeaterHendi.connect()).
		self.ser = serial.Serial()
		self.ser.port = port
		self.ser.baudrate = speed
		self.ser.timeout = 1.000

		powers = [500,  600,  700,  800,  900, 1000, 1100, 1200, 1300, 1400, 1500, 1600, 1700, 1800, 1900, 2000, 2100, 2200, 2300, 2400, 2500, 2600, 2700, 2800, 2900, 3000, 3100, 3200, 3300,3400, 3500]
		digits = [4092, 3840, 3700, 3440, 3300, 3060, 3000, 2930, 2860, 2780, 2640, 2510, 2430, 2360, 2190, 2130, 2030, 1980, 1880, 1730, 1630, 1560, 1480, 1400, 1240, 1150, 1150, 1080, 920, 740, 670]
		self.caps = {
			"pwr_watts_min"		: powers[0],
			"pwr_watts_max"		: powers[-1],
			"pwr_list"			: powers,
			"pwr_watts_step"	: 100,
			"pwr_digits_min"	: digits[0],
			"pwr_digits_max"	: digits[-1],
			"digit_list"		: digits
		}
		self.poly_w2d = np.polyfit(powers, digits, 5)
		self.poly_d2w = np.polyfit(digits, powers, 5)

	def is_connected(self):
		return self.ser.is_open

	def connect(self):
		if self.ser.is_open:
			return True
		try:
			self.ser.open()
		except serial.SerialException as e:
			print(f"HendiCtrl: connect failed: {e}")
			return False
		# The port itself is what "connected" means - a failed version query
		# (e.g. the device sitting in its ungraceful-disconnect lockout, see
		# docs/hendi_lockout_findings.md) doesn't make the link any less
		# open, so it's reported but doesn't fail connect() overall.
		try:
			self.sw_id = self.getSoftwareIdentifier()
			self.sw_ver = self.getSoftwareVersion()
			print(f"{self.sw_id}, F/W-Version: {self.sw_ver}")
		except HendiException:
			print("HendiCtrl: Hendi not found")
		return True

	@contextmanager
	def remote_open(self):
		try:
			self.remoteEnable(True)
			yield None
		finally:
			self.remoteEnable(False)

	def toDigits(self, power_watts):
		return np.polyval(self.poly_w2d, float(power_watts))

	def toWatts(self, power_digits):
		return np.polyval(self.poly_d2w, float(power_digits))

	def reset(self):
		self.ser.dtr = True
		self.ser.rts = False
		time.sleep(0.01)
		self.ser.dtr = True
		self.ser.rts = True

	def enter_bootloader(self):
		self.ser.dtr = True
		self.ser.rts = False
		time.sleep(0.01)
		self.ser.dtr = False
		self.ser.rts = True

	def firmware_update(self, filename):
		print("Start firmware update")
		self.enter_bootloader()
		self.ser.readline()

		self.ser.flushInput()
		with open(filename, 'r') as fp:
			line_count = 0
			while (True):
				s = fp.readline()
				if s == '':
					break
				data = s.encode()
				self.ser.write(data)
				ack = self.ser.read()
				sys.stdout.write(ack.decode('utf-8'))
				sys.stdout.flush()
				if line_count == 64:
					line_count = 0
					sys.stdout.write("\n")

				line_count += 1

		self.ser.readline()
		self.ser.flushInput()
		self.ser.write(b'q')
		self.ser.flushOutput()

	def __write(self, s):
		self.ser.flushInput()
		data = s.encode()
		if self.debug:
			print(f"HendiCtrl: -> {data!r}")
		self.ser.write(data + b'\r')

	def __read(self):
		# Read line echo
		echo = self.ser.readline()

		# Read answer line
		answer_raw = self.ser.readline()
		if self.debug:
			print(f"HendiCtrl: <- echo={echo!r} answer={answer_raw!r}")
		answer = answer_raw.decode('utf-8').replace('\r', '').replace('\n', '')
		if ':' in answer:
			result = answer.split(':')
		else:
			result = answer

		return result, echo, answer_raw

	def cmd(self, req):
		self.__write(req)
		rsp, echo, answer_raw = self.__read()
		if len(rsp) == 0:
			raise HendiException(f"Communication error: req={req!r} echo={echo!r} answer={answer_raw!r}")
		if "OK" not in rsp[0]:
			raise HendiException(f"Communication error: req={req!r} echo={echo!r} answer={answer_raw!r}")

		return rsp[1]
		
	def getCapabilties(self):
		return self.caps

	def getInfo(self):
		result = {
			"sw_id"			: self.getSoftwareIdentifier(),
			"sw_ver"		: self.getSoftwareVersion(),
			'state'			: self.getState(),
			'pwr_digits'	: self.getPowerDigits(),
			'Switch_state'	: self.getSwitch()
		}
		return result

	def _debugEnable(self, enable):
		return self.cmd("D" + str(int(enable)))

	def _isDebugEnable(self):
		return self.cmd("D?")

	def getSoftwareIdentifier(self):
		return self.cmd("I?")

	def getSoftwareVersion(self):
		return self.cmd("V?")

	def getState(self):
		return self.cmd("T?")

	def remoteEnable(self, enable):
		return self.cmd("R" + str(int(enable)))

	def isRemoteEnable(self):
		return '1' in self.cmd("R?")

	def setPowerDigits(self, power_digits):
		return self.cmd("P" + str(power_digits))

	def getPowerDigits(self):
		return self.cmd("P?")

	def setSwitch(self, enable):
		return self.cmd("S" + str(int(enable)))

	def getSwitch(self):
		return self.cmd("S?")

	def setPowerWatts(self, power_watts):
		pwr_digits_min = self.caps['pwr_digits_min']
		pwr_digits_max = self.caps['pwr_digits_max']
		max_digits = max(pwr_digits_max, pwr_digits_min)
		min_digits = min(pwr_digits_max, pwr_digits_min)
		pwr_digits = int(self.toDigits(power_watts))
		pwr_digits = int(min(max_digits, max(min_digits, pwr_digits)))
		return self.setPowerDigits(pwr_digits)

	def getPowerWatts(self):
		return int(self.toWatts(self.getPowerDigits()) + 0.5)

	def disconnect(self):
		if self.ser.is_open:
			try:
				self.remoteEnable(False)
			except HendiException:
				pass
			self.ser.close()
		self.sw_id = None
		self.sw_ver = None

	def close(self):
		self.disconnect()