import usb.core
import usb.util
import time
import datetime
import RPi.GPIO as GPIO
import Adafruit_GPIO.SPI as SPI
import Adafruit_SSD1306

from PIL import Image
from PIL import ImageDraw
from PIL import ImageFont

import subprocess


USB_IF = 0  # Interface
USB_TIMEOUT = 5  # Timeout in MS

SLEEP_TIME = 0.01

BTN_LEFT = 4
BTN_RIGHT = 8
BTN_DOWN = 2
BTN_UP = 1
BTN_STOP = 44  # Space
BTN_EXIT = 32  # ESC

USB_VENDOR = 0x045E  # Rii
USB_PRODUCT = 0x028E  # Mini Wireless Keyboard

# Raspberry Pi pin configuration:
RST = None  # on the PiOLED this pin isnt used
# Note the following are only used with SPI:
DC = 23
SPI_PORT = 0
SPI_DEVICE = 0

ForwardRight = 20
BackwardRight = 26

ForwardLeft = 18
BackwardLeft = 17


def forward():
    GPIO.output(ForwardLeft, GPIO.HIGH)
    GPIO.output(ForwardRight, GPIO.HIGH)
    time.sleep(SLEEP_TIME)


def backward():
    GPIO.output(BackwardRight, GPIO.HIGH)
    GPIO.output(BackwardLeft, GPIO.HIGH)
    time.sleep(SLEEP_TIME)


def turn_left():
    GPIO.output(ForwardLeft, GPIO.HIGH)
    GPIO.output(BackwardRight, GPIO.HIGH)
    time.sleep(SLEEP_TIME)


def turn_right():
    GPIO.output(ForwardRight, GPIO.HIGH)
    GPIO.output(BackwardLeft, GPIO.HIGH)
    time.sleep(SLEEP_TIME)


def clear():
    GPIO.output(ForwardLeft, GPIO.LOW)
    GPIO.output(ForwardRight, GPIO.LOW)
    GPIO.output(BackwardRight, GPIO.LOW)
    GPIO.output(BackwardLeft, GPIO.LOW)


def init_display():
    # create display instance
    disp = Adafruit_SSD1306.SSD1306_128_64(
        rst=RST, dc=DC, spi=SPI.SpiDev(SPI_PORT, SPI_DEVICE, max_speed_hz=8000000)
    )

    # Initialize library.
    disp.begin()

    # Clear display.
    disp.clear()
    disp.display()

    # Create blank image for drawing.
    # Make sure to create image with mode '1' for 1-bit color.
    width = disp.width
    height = disp.height
    image = Image.new("1", (width, height))

    # Get drawing object to draw on image.
    draw = ImageDraw.Draw(image)

    # Draw a black filled box to clear the image.
    draw.rectangle((0, 0, width, height), outline=0, fill=0)
    return disp, draw, image


def draw_stats(disp, draw, image):
    width = disp.width
    height = disp.height
    # Draw a black filled box to clear the image.
    draw.rectangle((0, 0, width, height), outline=0, fill=0)

    # Shell scripts for system monitoring from here : https://unix.stackexchange.com/questions/119126/command-to-display-memory-usage-disk-usage-and-cpu-load
    cmd = "hostname -i | cut -d' ' -f1"
    IP = subprocess.check_output(cmd, shell=True)
    cmd = "top -bn1 | grep load | awk '{printf \"CPU Load: %.2f\", $(NF-2)}'"
    CPU = subprocess.check_output(cmd, shell=True)
    cmd = "free -m | awk 'NR==2{printf \"Mem: %s/%sMB %.2f%%\", $3,$2,$3*100/$2 }'"
    MemUsage = subprocess.check_output(cmd, shell=True)
    cmd = 'df -h | awk \'$NF=="/"{printf "Disk: %d/%dGB %s", $3,$2,$5}\''
    Disk = subprocess.check_output(cmd, shell=True)

    # Write two lines of text.

    draw.text((x, top), "IP: " + str(IP), font=font, fill=255)
    # draw.text((x, top + 8), str(CPU), font=font, fill=255)
    draw.text((x, top + 19), str(MemUsage), font=font, fill=255)
    draw.text((x, top + 38), str(Disk), font=font, fill=255)
    draw.text((x, top + 47), str(datetime.datetime.now().time()), font=font, fill=255)

    # Display image.
    disp.image(image)
    disp.display()
    time.sleep(0.5)


GPIO.setmode(GPIO.BCM)
GPIO.setup(ForwardRight, GPIO.OUT)
GPIO.setup(BackwardRight, GPIO.OUT)
GPIO.setup(ForwardLeft, GPIO.OUT)
GPIO.setup(BackwardLeft, GPIO.OUT)

disp, draw, image = init_display()
width = disp.width
height = disp.height
# First define some constants to allow easy resizing of shapes.
padding = -2
top = padding
bottom = height - padding
# Move left to right keeping track of the current x position for drawing shapes.
x = 0
# Load default font.
font = ImageFont.load_default()

# main logic starts
dev = usb.core.find(idVendor=USB_VENDOR, idProduct=USB_PRODUCT)

endpoint = dev[0][(0, 0)][0]

if dev.is_kernel_driver_active(USB_IF) is True:
    dev.detach_kernel_driver(USB_IF)

usb.util.claim_interface(dev, USB_IF)

draw_stats(disp, draw, image)

while True:
    control = None

    try:
        control = dev.read(
            endpoint.bEndpointAddress, endpoint.wMaxPacketSize, USB_TIMEOUT
        )
        if control[0] == 0:
            # print('clear')
            clear()
    except:
        pass
    if control is not None:
        if control[2] == BTN_DOWN:
            print("-")
            backward()

        if control[2] == BTN_UP:
            print("+")
            forward()

        if control[2] == BTN_LEFT:
            print("<-")
            turn_left()

        if control[2] == BTN_RIGHT:
            print("->")
            turn_right()

        if control[2] == BTN_STOP:
            print("stop")

        if control[2] == BTN_EXIT:
            clear()
            exit()
        if control[3] == 2:
            print("fresh stats")
            draw_stats(disp, draw, image)
            continue
    else:
        # Let CTRL+C actually exit
        clear()
    time.sleep(SLEEP_TIME)