#!/usr/bin/env python3
import argparse
import dataclasses
import lgpio
import sys
from typing import Dict
import time

# From test C program using wiringPi
# RELAY_PINS = [4, 17, 27, 18, 22, 23, 24, 25]
# Physical pins
# PINS = [7, 11, 13, 15, 12, 16, 18, 22]
# From Banana Pi Zero M2 pinout
RELAY_PINS = ["CON2-P07", "CON2-P11", "CON2-P13", "CON2-P15", "CON2-P12", "CON2-P16", "CON2-P18", "CON2-P22"]
# con2-p22 ?
# p18: waterfall
# p16: filter pump, salt cell

@dataclasses.dataclass
class Relay:
	name: str
	pin: str

RELAY_ON = 0
RELAY_OFF = 1

RELAYS = [
	Relay("Unknown", "CON2-P22"),
	Relay("Waterfall", "CON2-P18"),
	Relay("Filter", "CON2-P16"),
	Relay("Unknown 2", "CON2-P12"),
]
DEVICE = "/dev/ttyUSB0"

def alloff(args, handle, pin_to_line):
	"Turn off all relays"
	for relay in RELAYS:
		line = pin_to_line[relay.pin]
		lgpio.gpio_claim_output(handle, line, level=RELAY_OFF)

def on(args, handle, pin_to_line):
	"Turn on a single relay"
	name_to_pin = {relay.name.lower(): relay.pin for relay in RELAYS}
	pin = name_to_pin[args.name.lower()]
	line = pin_to_line[pin]
	lgpio.gpio_claim_output(handle, line, level=RELAY_ON)
	print(f"{args.name} now on")
	
def status(args, handle, pin_to_line):
	pin_to_name = {relay.pin: relay.name for relay in RELAYS}
	for pin, name in pin_to_name.items():
		line = pin_to_line[pin]
		okay_status, io_number, flags, pin, user = lgpio.gpio_get_line_info(handle, line)
		# value 2 == "Kernel: Output"
		print(f"{pin_to_name[pin]}: {flags} ({okay_status})")

def map_pin_to_line(handle):
	start_pin, end_pin, name, driver = lgpio.gpio_get_chip_info(handle)
	pin_to_line = {}
	for x in range(start_pin, end_pin):
		okay_status, io_number, flags, pin, user = lgpio.gpio_get_line_info(handle, x)
		if okay_status != 0:
			print(f"Error getting pin {x}: status {okay_status}")
			continue
		pin_to_line[pin] = x
	return pin_to_line


def main():
	parser = argparse.ArgumentParser()
	subparsers = parser.add_subparsers(required=True, help="the sub-command to execute")

	parser_status = subparsers.add_parser("status", help="Show status of the relays")
	parser_status.set_defaults(func=status)

	parser_alloff = subparsers.add_parser("alloff", help="Turn off all relays")
	parser_alloff.set_defaults(func=alloff)

	parser_on = subparsers.add_parser("on", help="Turn on a relay")
	parser_on.add_argument("name", choices=[relay.name for relay in RELAYS])
	parser_on.set_defaults(func=on)
	args = parser.parse_args()

	handle = lgpio.gpiochip_open(0) # open /dev/gpiochip0
	if handle < 0:
		print("Error on open")
		return -1

	pin_to_line = map_pin_to_line(handle)
		
	try:
		args.func(args, handle, pin_to_line)
	finally:
		lgpio.gpiochip_close(handle)

	return 0

if __name__ == "__main__":
	sys.exit(main())