#!/usr/bin/env python3

import os
import threading
import time

KEY1=21
KEY2=20
KEY3=16
LEFT=5
RIGHT=26
UP=6
DOWN=19
PRESS=13

KEYS = [KEY1,KEY2,KEY3,UP,DOWN,LEFT,RIGHT,PRESS]

export = os.open('/sys/class/gpio/export', os.O_WRONLY | os.O_NONBLOCK)

for gpio in KEYS:
    if not os.path.exists('/sys/class/gpio/gpio%d' % gpio):
        os.write(export,  bytes(str(gpio), 'ascii'))
    direction = os.open('/sys/class/gpio/gpio%d/direction' % gpio, os.O_WRONLY | os.O_NONBLOCK)
    os.write(direction, bytes('high', 'ascii'))
    os.close(direction)
os.close(export)

# Reset USB Gadget
os.system('echo > /sys/kernel/config/usb_gadget/procon/UDC')
os.system('ls /sys/class/udc > /sys/kernel/config/usb_gadget/procon/UDC')

time.sleep(0.5)

gadget = os.open('/dev/hidg0', os.O_RDWR | os.O_NONBLOCK)

counter = 0
mac_addr = '00005e00535e'
initial_input = '81008000f8d77a22c87b0c'

def countup():
    global counter
    while True:
        counter = (counter + 3) % 256
        time.sleep(0.03)

def response(code, cmd, data):
    buf = bytearray([code, cmd])
    buf.extend(data)
    buf.extend(bytearray(64-len(buf)))
    try:
        os.write(gadget, buf)
    except BlockingIOError:
        pass
    except:
        os._exit(1)

def uart_response(code, subcmd, data):
    buf = bytearray.fromhex(initial_input)
    buf.extend([code, subcmd])
    buf.extend(data)

    response(0x21, counter, buf)

def spi_response(addr, data):
    buf = bytearray(addr)
    buf.extend([0x00, 0x00])
    buf.append(len(data))
    buf.extend(data)
    
    uart_response(0x90, 0x10, buf)

def input_response():
    while True:
        buf = bytearray.fromhex(initial_input)
        inputs = {}
        for gpio in KEYS:
            with open('/sys/class/gpio/gpio%d/value' % gpio, 'r') as f:
                inputs[gpio] = f.read().startswith('0')
        buf[1] = inputs[KEY1] << 3 | inputs[KEY2] << 2
        buf[2] |= inputs[KEY3] << 4 | inputs[PRESS] << 1
        buf[3] = inputs[LEFT] << 3 | inputs[RIGHT] << 2 | inputs[UP] << 1 | inputs[DOWN]
        
        response(0x30, counter, buf)
        time.sleep(0.03)

def simulate_procon():
    while True:
        try:
            data = os.read(gadget, 128)
            if data[0] == 0x80:
                if data[1] == 0x01:
                    response(0x81, data[1], bytes.fromhex('0003' + mac_addr))
                elif data[1] == 0x02:
                    response(0x81, data[1], [])
                elif data[1] == 0x04:
                    threading.Thread(target=input_response).start()
                else:
                    print('>>>', data.hex())
            elif data[0] == 0x01 and len(data) > 16:
                if data[10] == 0x01: # Bluetooth manual pairing
                    uart_response(0x81, data[10], [0x03])
                elif data[10] == 0x02: # Request device info
                    uart_response(0x82, data[10], bytes.fromhex('03480302' + mac_addr[::-1] + '0301'))
                elif data[10] == 0x03 or data[10] == 0x08 or data[10] == 0x30 or data[10] == 0x38 or data[10] == 0x40 or data[10] == 0x48:                
                    uart_response(0x80, data[10], [])
                elif data[10] == 0x04: # Trigger buttons elapsed time
                    uart_response(0x83, data[10], [])
                elif data[10] == 0x21: # Set NFC/IR MCU configuration
                    uart_response(0xa0, data[10], bytes.fromhex('0100ff0003000501'))
                elif data[10] == 0x10:
                    if data[11:13] == b'\x00\x60': # Serial number
                        spi_response(data[11:13], bytes.fromhex('ffffffffffffffffffffffffffffffff'))
                    elif data[11:13] == b'\x50\x60': # Controller Color
                        spi_response(data[11:13], bytes.fromhex('bc1142 75a928 ffffff ffffff ff')) # Raspberry Color
                    elif data[11:13] == b'\x80\x60': # Factory Sensor and Stick device parameters
                        spi_response(data[11:13], bytes.fromhex('50fd0000c60f0f30619630f3d41454411554c7799c333663'))
                    elif data[11:13] == b'\x98\x60': # Factory Stick device parameters 2
                        spi_response(data[11:13], bytes.fromhex('0f30619630f3d41454411554c7799c333663'))
                    elif data[11:13] == b'\x3d\x60': # Factory configuration & calibration 2
                        spi_response(data[11:13], bytes.fromhex('ba156211b87f29065bffe77e0e36569e8560ff323232ffffff'))
                    elif data[11:13] == b'\x10\x80': # User Analog sticks calibration
                        spi_response(data[11:13], bytes.fromhex('ffffffffffffffffffffffffffffffffffffffffffffb2a1'))
                    elif data[11:13] == b'\x28\x80': # User 6-Axis Motion Sensor calibration
                        spi_response(data[11:13], bytes.fromhex('beff3e00f001004000400040fefffeff0800e73be73be73b'))
                    else:
                        print("Unknown SPI address:", data[11:13].hex())
                else:
                    print('>>> [UART]', data.hex())
            elif data[0] == 0x10 and len(data) == 10:
                pass
            else:
                print('>>>', data.hex())
        except BlockingIOError:
            pass
        except:
            os._exit(1)

threading.Thread(target=simulate_procon).start()
threading.Thread(target=countup).start()