import hid
import time
import struct

DEFAULT_HID_TIMEOUT = 10 # milliseconds to wait for a HID message

BIGSCREEN_VID = 0x35BD
BEYOND_PID = 0x0101

HMD_GET_SW_VER = ord('*')
HMD_GET_BOARD_SERIAL = ord('%')
HMD_GET_VXR_FIRMWARE = ord('N')
HMD_GET_OLED_ID = ord('^')
HMD_GET_USAGE_TIMER = ord('Z')
HMD_ERROR = ord('E')

# "wait_for_response" is a helper function that waits for a specific message type
# the HID response packet (aka the HID report) has the message type as the first byte
# send the desired message types in a list of ints, any message that 
# matches one of the types in the list will be returned
# if timeout, returns empty bytes
#
# params:
# hmd_device -      A hid.Device object that is already connected to the HMD
# message_types -   List of ints, each int is a message type. If any message
#                   is received that matches one of the types, this function
#                   will return with that response packet
# timeout_ms -      An int with the desired timeout in milliseconds. Default 1000 
#
# Returns:  A bytes object. If timed out, will be empty bytes (b''). Otherwise
#           it will have the first message that matched one of the types given.

def wait_for_response(hmd_device: hid.Device, message_types:list, timeout_ms: int = 1000) -> bytes:
    start_time = time.monotonic_ns()
    
    while( (start_time + (timeout_ms*1000000)) > time.monotonic_ns()):
        bytesout = hmd_device.read(65, timeout=DEFAULT_HID_TIMEOUT)
        if(len(bytesout) > 0):
            if(bytesout[0] in message_types):
                return bytesout
    
    return b''

# Gets the microcontroller firmware version as a string. If not found, returns the string "error"
def get_hmd_sw_ver(hmd_device: hid.Device) -> str:
    hmd_device.send_feature_report(bytes([0, HMD_GET_SW_VER]))

    response = wait_for_response(hmd_device, [HMD_GET_SW_VER, HMD_ERROR])
    if(len(response) == 0 or response[0] == HMD_ERROR):
        return "error"
    
    return response[1:].rstrip(b'\x00').decode('ascii')

def get_hmd_board_serial(hmd_device: hid.Device) -> str:
    hmd_device.send_feature_report(bytes([0, HMD_GET_BOARD_SERIAL]))
    response = wait_for_response(hmd_device, [HMD_GET_BOARD_SERIAL, HMD_ERROR])
    if(len(response) == 0 or response[0] == HMD_ERROR):
        return "error"
    
    return response[1:].rstrip(b'\x00').decode('ascii')

# Gets the DisplayPort bridge IC firmware version as a string. If not found, returns the string "error"
def get_hmd_vxr_ver(hmd_device: hid.Device) -> str:
    hmd_device.send_feature_report(bytes([0, HMD_GET_VXR_FIRMWARE]))
    response = wait_for_response(hmd_device, [HMD_GET_VXR_FIRMWARE, HMD_ERROR])
    if(len(response) == 0 or response[0] == HMD_ERROR):
        return "error"
    
    return response[1:].rstrip(b'\x00').decode('ascii')

# returns two strings in a tuple: left display then right display
def get_hmd_oled_id(hmd_device: hid.Device) -> tuple:

    hmd_device.send_feature_report(bytes([0, HMD_GET_OLED_ID, 0]))
    response = wait_for_response(hmd_device, [HMD_GET_OLED_ID, HMD_ERROR])
    if(len(response) == 0 or response[0] == HMD_ERROR or response[1] == 0):
        left_id = "error"
    else:
        left_id = response[2:].rstrip(b'\x00').decode('ascii')
    hmd_device.send_feature_report(bytes([0, HMD_GET_OLED_ID, 1]))
    response = wait_for_response(hmd_device, [HMD_GET_OLED_ID, HMD_ERROR])
    if(len(response) == 0 or response[0] == HMD_ERROR or response[1] == 0):
        right_id = "error"
    else:
        right_id = response[2:].rstrip(b'\x00').decode('ascii')

    return (left_id, right_id)


def get_usage_timer(hmd_device: hid.Device, timer_num: int) -> int:
    hmd_device.send_feature_report(bytes([0, HMD_GET_USAGE_TIMER, timer_num]))
    response = wait_for_response(hmd_device, [HMD_GET_USAGE_TIMER, HMD_ERROR])
    if(len(response) == 0 or response[0] == HMD_ERROR or response[1] == 0):
        timer_val = -1
    else:
        timer_val = struct.unpack('<I',response[1:5])
    return timer_val

if __name__ == '__main__':
    myhmd = hid.Device(vid=BIGSCREEN_VID, pid = BEYOND_PID)

    print('HMD Software Version: ', end='')
    print(get_hmd_sw_ver(myhmd))

    print('VXR Software Version: ', end='')
    print(get_hmd_vxr_ver(myhmd))

    print('Board Serial Number: ', end='')
    print(get_hmd_board_serial(myhmd))

    (left_serial, right_serial) = get_hmd_oled_id(myhmd)

    print('Left OLED ID: ',end='')
    print(left_serial)
    print('Right OLED ID: ', end='')
    print(right_serial)

    total_on_time = get_usage_timer(myhmd, 0)
    display_on_time = get_usage_timer(myhmd, 1)
    longest_display_on_time = get_usage_timer(myhmd, 2)
    display_on_time_right = get_usage_timer(myhmd, 3)

    print('Total on time: {}'.format(total_on_time))
    print('Display on time (left): {}'.format(display_on_time))
    print('Display on time (right): {}'.format(display_on_time_right))
    print('Longest display on time: {}'.format(longest_display_on_time))