#!/usr/bin/env python3

# '''Enable Auto-Shutdown Protection Function '''
import os
import time
import smbus2
import logging
import argparse
from ina219 import INA219,DeviceRangeError


PROG_NAME = "UPS Control"
PROG_VERSION = 'v' + open(os.path.join(os.path.dirname(os.path.realpath(__file__)), "VERSION")).readline()
PROG_DESCRIPTION = ""
PROG_EPILOG = ""

# Define I2C bus
DEVICE_BUS = 1
# Define device i2c slave address.
DEVICE_ADDR = 0x17
# Set the threshold of UPS automatic power-off to prevent damage caused by battery over-discharge, unit: mV.
PROTECT_VOLT = 3500

logger = logging.getLogger("upsctl")
logging.basicConfig(level=logging.INFO)


def check_ups_status():
    # Instance INA219 and getting information from it.
    ina_supply = INA219(0.00725, busnum=DEVICE_BUS, address=0x40)
    ina_supply.configure()
    supply_voltage = ina_supply.voltage()
    supply_current = ina_supply.current()
    supply_power = ina_supply.power()

    logger.debug(" - Current information of the detected Raspberry Pi")
    logger.debug("Raspberry Pi Supply Voltage: %.3f V" % supply_voltage)
    logger.debug("Raspberry Pi Current Current Consumption: %.3f mA" % supply_current)
    logger.debug("Raspberry Pi Current Power Consumption: %.3f mW" % supply_power)

    # Batteries information
    ina_batt = INA219(0.005, busnum=DEVICE_BUS, address=0x45)
    ina_batt.configure()
    batt_voltage = ina_batt.voltage()
    batt_current = ina_batt.current()
    batt_power = ina_batt.power()

    logger.debug(" - Batteries information")
    logger.debug("Voltage of Batteries: %.3f V" % batt_voltage)
    try:
        if batt_current > 0:
            logger.debug("Battery Current (Charging) Rate: %.3f mA"% batt_current)
            logger.debug("Current Battery Power Supplement: %.3f mW"% batt_power)
        else:
            logger.debug("Battery Current (discharge) Rate: %.3f mA"% batt_current)
            logger.debug("Current Battery Power Consumption: %.3f mW"% batt_power)
    except DeviceRangeError:
        logger.error('Battery power is too high.')

    # Raspberry Pi Communicates with MCU via i2c protocol.
    bus = smbus2.SMBus(DEVICE_BUS)

    aReceiveBuf = []
    aReceiveBuf.append(0x00)

    # Read register and add the data to the list: aReceiveBuf
    for i in range(1, 255):
        aReceiveBuf.append(bus.read_byte_data(DEVICE_ADDR, i))

    # UID0 = "%08X" % (aReceiveBuf[243] << 24 | aReceiveBuf[242] << 16 | aReceiveBuf[241] << 8 | aReceiveBuf[240])
    # UID1 = "%08X" % (aReceiveBuf[247] << 24 | aReceiveBuf[246] << 16 | aReceiveBuf[245] << 8 | aReceiveBuf[244])
    # UID2 = "%08X" % (aReceiveBuf[251] << 24 | aReceiveBuf[250] << 16 | aReceiveBuf[249] << 8 | aReceiveBuf[248])
    # logger.info('UID:' + UID0 + '/' + UID1 + '/' + UID2)

    logger.debug(" - Current state")
    if (aReceiveBuf[8] << 8 | aReceiveBuf[7]) > 4000:
        logger.debug('Currently charging via Type C Port')
    elif (aReceiveBuf[10] << 8 | aReceiveBuf[9]) > 4000:
        logger.debug('Currently charging via Micro USB Port')
    else:
        logger.warning('Currently not charging.')
        # Consider shutting down to save data or send notifications
        if (str(batt_voltage)) == ("0.0"):
            logger.error("Bad battery voltage value")
        if (str(batt_voltage)) != ("0.0"):
            if ((batt_voltage * 1000) < (PROTECT_VOLT + 200)):
                logger.warning('The battery is going to dead! Ready to shut down!')
                # It will cut off power when initialized shutdown sequence.
                bus.write_byte_data(DEVICE_ADDR, 24, 240)
                os.system("sync && halt")
                while True:
                    time.sleep(10)


def init_ups():
    logger.info("Setting initial parameters for the UPS")
    # Raspberry Pi Communicates with MCU via i2c protocol.
    bus = smbus2.SMBus(DEVICE_BUS)

    # Enable Back-to-AC fucntion.
    # Enable: write 1 to register 0x19 == 25
    # Disable: write 0 to register 0x19 == 25
    bus.write_byte_data(DEVICE_ADDR, 25, 1)
    logger.info("Back-to-AC function enabled")

    # Reset Protect voltage
    bus.write_byte_data(DEVICE_ADDR, 17, PROTECT_VOLT & 0xFF)
    bus.write_byte_data(DEVICE_ADDR, 18, (PROTECT_VOLT >> 8)& 0xFF)
    logger.info(f"Successfully set the protection voltage to: {PROTECT_VOLT} mV")


def print_status():
    # Instance INA219 and getting information from it.
    ina_supply = INA219(0.00725, busnum=DEVICE_BUS, address=0x40)
    ina_supply.configure()
    supply_voltage = ina_supply.voltage()
    supply_current = ina_supply.current()
    supply_power = ina_supply.power()
    print("-" * 60)
    print("----- Current information of the detected Raspberry Pi -----")
    print("-" * 60)
    print("Raspberry Pi Supply Voltage: %.3f V" % supply_voltage)
    print("Raspberry Pi Current Current Consumption: %.3f mA" % supply_current)
    print("Raspberry Pi Current Power Consumption: %.3f mW" % supply_power)
    print("-" * 60)

    # Batteries information
    ina_batt = INA219(0.005, busnum=DEVICE_BUS, address=0x45)
    ina_batt.configure()
    batt_voltage = ina_batt.voltage()
    batt_current = ina_batt.current()
    batt_power = ina_batt.power()
    print("------------------ Batteries information ------------------")
    print("-" * 60)
    print("Voltage of Batteries: %.3f V" % batt_voltage)
    try:
        if batt_current > 0:
            print("Battery Current (Charging) Rate: %.3f mA" % batt_current)
            print("Current Battery Power Supplement: %.3f mW" % batt_power)
        else:
            print("Battery Current (discharge) Rate: %.3f mA" % batt_current)
            print("Current Battery Power Consumption: %.3f mW" % batt_power)
            print("-" * 60)
    except DeviceRangeError:
        print("-" * 60)
        print('Battery power is too high.')

    # Raspberry Pi Communicates with MCU via i2c protocol.
    bus = smbus2.SMBus(DEVICE_BUS)

    aReceiveBuf = []
    aReceiveBuf.append(0x00)

    # Read register and add the data to the list: aReceiveBuf
    for i in range(1, 255):
        aReceiveBuf.append(bus.read_byte_data(DEVICE_ADDR, i))

    # UID0 = "%08X" % (aReceiveBuf[243] << 24 | aReceiveBuf[242] << 16 | aReceiveBuf[241] << 8 | aReceiveBuf[240])
    # UID1 = "%08X" % (aReceiveBuf[247] << 24 | aReceiveBuf[246] << 16 | aReceiveBuf[245] << 8 | aReceiveBuf[244])
    # UID2 = "%08X" % (aReceiveBuf[251] << 24 | aReceiveBuf[250] << 16 | aReceiveBuf[249] << 8 | aReceiveBuf[248])
    # print('UID:' + UID0 + '/' + UID1 + '/' + UID2)

    if (aReceiveBuf[8] << 8 | aReceiveBuf[7]) > 4000:
        print('-' * 60)
        print('Currently charging via Type C Port.')
    elif (aReceiveBuf[10] << 8 | aReceiveBuf[9]) > 4000:
        print('-' * 60)
        print('Currently charging via Micro USB Port.')
    else:
        print('-' * 60)
        print('Currently not charging.')
        if (str(batt_voltage)) == ("0.0"):
            print("Bad battery voltage value")
        if (str(batt_voltage)) != ("0.0"):
            if ((batt_voltage * 1000) < (PROTECT_VOLT + 200)):
                print('-' * 60)
                print('The battery is going to dead!')


def main():
    arg_parse = argparse.ArgumentParser(description='UPS Control Script')
    arg_parse.add_argument('-s', '--status', action='store_true', help='Print the status of the UPS')
    arg_parse.add_argument('-i', '--init', action='store_true', help='Set initial parameters for the UPS')
    arg_parse.add_argument('-c', '--check', action='store_true', help='Check the status of the UPS')
    args = arg_parse.parse_args()

    if args.check:
        check_ups_status()
    elif args.init:
        init_ups()
    elif args.status:
        print_status()
    else:
        arg_parse.print_help()


if __name__ == "__main__":
    main()