#! /usr/bin/python3

########################################################
# pi-top battery level gauge (pt-pulse-battery)        #
# using the pi-topPULSE and getting battery state info #
# from the pt-device-manager                           #
# Original by Jez - @JezShed for Unicorn hat           #
# Code to interact with device manager by pi-top       #
# Modified by @rricharz for pi-topPULSE                #
# version 1.1                                          #
# Note:                                                #
# This script requires either the latest pi-topOS      #
# with the pt-device-manager, or Raspbian Stretch      #
# with the latest pt-device-manager installed          #
########################################################

import colorsys
import time
import subprocess
import sys
from ptpulse import ledmatrix
import zmq
from ptcommon.ptdm_message import Message

zmq_socket = None


def connect_to_socket():
    global zmq_socket

    zmq_context_send = zmq.Context()
    zmq_socket = zmq_context_send.socket(zmq.REQ)
    zmq_socket.sndtimeo = 1000
    zmq_socket.rcvtimeo = 1000
    zmq_socket.connect("tcp://127.0.0.1:3782")


def send_request(message_request_id, parameters):

    message = Message.from_parts(message_request_id, parameters)
    zmq_socket.send_string(message.to_string())

    response_string = zmq_socket.recv_string()
    return Message.from_string(response_string)


def cleanup():
    if (zmq_socket is not None):
        zmq_socket.close(0)


def get_battery_state_message():
    message = Message.from_parts(Message.REQ_GET_BATTERY_STATE, [])
    zmq_socket.send_string(message.to_string())

    response_string = zmq_socket.recv_string()
    return Message.from_string(response_string)


def drawbattery():  # Draw the naked battery
    for y in range(0, 6):
        ledmatrix.set_pixel(1, y, 64, 64, 255)
        ledmatrix.set_pixel(5, y, 64, 64, 255)
    for x in range(2, 5):
        ledmatrix.set_pixel(x, 0, 64, 64, 255)
        ledmatrix.set_pixel(x, 6, 192, 192, 192)
    ledmatrix.show()


def show_capacity(message):
    # Plot the battery level on the pi-topPULSE
    if message.message_id() == Message.RSP_GET_BATTERY_STATE:
        if message.validate_parameters([int, int, int, int]):
            str_charging_state, str_capacity, str_time_remaining, str_wattage = message.parameters()
            # print("Charging State: " + str_charging_state)
            # print("Capacity: " + str_capacity)
            charging_state = int(str_charging_state)
            capacity = int(str_capacity)
        else:
            print("Invalid parameter types received from device manager.")
            time.sleep(5)
            return 2
    else:
        print("Unable to get valid battery information.")
        time.sleep(5)
        return 1

    r = 0
    g = 0
    b = 0
    if charging_state == 0:
        if capacity < 11:
            r = 255
        else:
            g = 255
    elif charging_state == 1:
        r = 255
        g = 225

    cap = int(capacity / 20) + 1
    if cap < 0:
        cap = 0
    if cap > 5:
        cap = 5

    if cap > 0:
        for y in range(1, cap+1):
            ledmatrix.set_pixel(2, y, r, g, b)
            ledmatrix.set_pixel(3, y, r, g, b)
            ledmatrix.set_pixel(4, y, r, g, b)
    if cap == 0:
        cap = 1
    if cap < 6:
        if (capacity < 50) and (charging_state == 0):
            # blinking warning
            for i in range(1, 3):
                for y in range(cap+1, 6):
                    ledmatrix.set_pixel(2, y, 0, 0, 0)
                    ledmatrix.set_pixel(3, y, 0, 0, 0)
                    ledmatrix.set_pixel(4, y, 0, 0, 0)
                ledmatrix.show()
                time.sleep(0.4)
                for y in range(cap+1, 6):
                    ledmatrix.set_pixel(2, y, 255, 0, 0)
                    ledmatrix.set_pixel(3, y, 255, 0, 0)
                    ledmatrix.set_pixel(4, y, 255, 0, 0)
                ledmatrix.show()
                time.sleep(0.4)

        else:
            for y in range(cap+1, 6):
                ledmatrix.set_pixel(2, y, 0, 0, 0)
                ledmatrix.set_pixel(3, y, 0, 0, 0)
                ledmatrix.set_pixel(4, y, 0, 0, 0)
            ledmatrix.show()
            time.sleep(5)
    return 0


############################
# Main program starts here #
############################
ledmatrix.rotation(0)
ledmatrix.clear()          # Clear the display
drawbattery()               # Draw the battery

while True:
    try:
        connect_to_socket()
        message = get_battery_state_message()
        show_capacity(message)
        cleanup()

    except Exception as e:
        print("Error getting battery info: " + str(e))
        time.sleep(5)
