from enum import IntEnum
from pyblake2 import blake2s
from pyLoraRFM9x import LoRa, ModemConfig
import time, toml, math, struct

HASH_LENGTH = 8


class MessageType(IntEnum):
	DeviceStatus = 1
	SensorStatus = 2
	Time = 2


class MultiNode:
	sensor_type_table = {1: "V", 2: "mA"}

	def __init__(self):
		with open("Config.toml", "r") as config_file:
			config = toml.loads(config_file.read())
			
		self.server_address = config["server"]["address"]
		self.server_secret_key = config["server"]["secret_key"]
		# How often to send the time
		self.time_interval = config["server"]["time_interval"]
		self.devices = {}
		self.last_time_message = time.time()

		for node_id_str in config["node"]:
			node_id = int(node_id_str, 16)
			self.devices[node_id] = config["node"][node_id_str]

		self.lora = LoRa(
		    0,  # SPI channel
		    25,  # Interrupt pin
		    255,  # Node ID
		    reset_pin=22,
		    modem_config=ModemConfig.Bw125Cr45Sf128,
		    tx_power=14,
		    freq=868,
		    acks=False)  # , receive_all=True)

		self.lora.cad_timeout = 1
		self.lora.on_recv = self.process_packet
		self.lora.set_mode_rx()

	def loop(self):
		if time.time() - self.last_time_message > self.time_interval:
			self.send_packet(0xFFFF, int(MessageType.Time).to_bytes(1, "little"))
			self.last_time_message += self.time_interval
			#print("Sent time")

	def decode_packet(self, device, data):
		packet_type = data[0]

		# match packet_type:
		#	case MessageType.DeviceStatus:

		if packet_type == MessageType.DeviceStatus:
			device["status"] = {"battery": struct.unpack('<f', data[1:5])[0], "temperature": struct.unpack('<f', data[5:10])[0]}
			return device["status"]

		if packet_type == MessageType.SensorStatus:
			channels_raw = struct.unpack('<H', data[1:3])[0]
			channels = []
			for i in range(16):
				if (channels_raw >> i) & 1:
					channels.append(i)

			sensor_data = []
			for i in range(len(channels)):
				offset = i * 6
				sensor_data.append({"channel": channels[i], "type": data[3 + offset], "pin": data[4 + offset], "value": struct.unpack('<f', data[5 + offset:9 + offset])[0]})
			device["sensors"] = sensor_data
			return sensor_data

	def process_packet(self, payload):
		rx_id = int.from_bytes(payload.message[0:2], byteorder="little")
		tx_id = int.from_bytes(payload.message[2:4], byteorder="little")
		msg_id = int.from_bytes(payload.message[4:8], byteorder="little")
		length = payload.message[8]
		data = payload.message[9:9 + length]
		data_hash = payload.message[9 + length:9 + length + HASH_LENGTH]

		if len(payload.message) != length + 9 + HASH_LENGTH:
			print(f"Invalid length! Expected {length + 9 + HASH_LENGTH} actual {len(payload.message)}")
			return

		if not tx_id in self.devices:
			with open("Config.toml", "r") as config_file:
				config = toml.loads(config_file.read())
				for node_id_str in config["node"]:
					node_id = int(node_id_str, 16)
					if not node_id in self.devices:
						self.devices[node_id] = config["node"][node_id_str]
			
			if not tx_id in self.devices:
				print(f"Error: Unrecognized device with ID {tx_id}")
				return

		hash_function = blake2s(key=self.devices[tx_id]["secret_key"].to_bytes(8, "little"), digest_size=8)
		hash_function.update(payload.message[:-HASH_LENGTH])

		if hash_function.digest() != data_hash:
			print(f"Hash doesn't match! Expected {hash_function.digest()} got {data_hash}")
			return
		#	self.devices[tx_id] = {}

		if "last_message_id" in self.devices[tx_id]:
			if msg_id <= self.devices[tx_id]["last_message_id"]:
				print(f'Error, expected an ID larger than {self.devices[tx_id]["last_message_id"]} but got {msg_id}')
				return
		
		self.devices[tx_id]["last_message_id"] = msg_id
		self.decode_packet(self.devices[tx_id], data)
		# RSSI cahnges from -13 to -27 on change from +20 to +2 dBm TX setting
		# print(f"{tx_id} #{msg_id}: {self.decode_packet(self.devices[rx_id], data)} V, {payload.rssi} dB(?) RSSI, {payload.snr} dB(?) SNR {(time.clock_gettime_ns(0)) / 1e9}")
		# print(f"{tx_id} #{msg_id}: {data.hex()} {self.decode_packet(data)}, {payload.rssi} dB(?) RSSI, {payload.snr} dB(?) SNR {(time.clock_gettime_ns(0)) / 1e9}")
		self.print_data()

	def send_packet(self, target: int, data):
		assert len(data) < 256

		payload = []
		payload.extend(target.to_bytes(2, "little"))
		payload.extend(self.server_address.to_bytes(2, "little"))
		# TODO: is this the best idea? Clock change would affect it badly
		payload.extend(int(time.time()).to_bytes(4, "little"))
		payload.append(len(data))
		payload.extend(data)

		hash_function = blake2s(key=self.server_secret_key.to_bytes(8, "little"), digest_size=8)
		hash_function.update(bytearray(payload))
		payload.extend(hash_function.digest())
		#print(payload)

		self.lora.send(payload, 255)
		self.lora.set_mode_rx()

	def print_data(self):
		for device_id in self.devices:
			device = self.devices[device_id]
			if "last_message_id" in device:
				print(f'Node {device_id} @ ID {device["last_message_id"]}:')

				if "status" in device:
					print(f'\t{device["status"]["battery"]:.3f} V {device["status"]["temperature"]:.1f} °C')

				if "sensors" in device:
					for sensor in device["sensors"]:
						if sensor["value"] == math.nan:
							print(f'\tCH {sensor["channel"]} on Pin {sensor["pin"]}: ERROR')
						else:
							print(f'\tCH {sensor["channel"]} on Pin {sensor["pin"]}: {sensor["value"]:.4f} {self.sensor_type_table[sensor["type"]]}')
				
				print()


if __name__ == "__main__":
	multinode = MultiNode()
	while True:
		multinode.loop()