musicmouse/espmusicmouse/host_driver/host_driver.py

import asyncio
from enum import Enum
import struct

from led_cmds import *

MAGIC_TOKEN_HOST_TO_FW = 0x1d6379e3
MAGIC_TOKEN_FW_TO_HOST = 0x10c65631


class MessageFwToHost(Enum):
    RFID_TOKEN_READ = 0
    ROTARY_ENCODER = 1
    TOUCH_BUTTON_PRESS = 2
    TOUCH_BUTTON_RELEASE = 3


class TouchButton(Enum):
    LEFT_FOOT = 0
    RIGHT_FOOT = 1
    LEFT_EAR = 2
    RIGHT_EAR = 3


led_ring_effect_to_message_id = {
    EffectStaticConfig: 0,
    EffectAlexaSwipeConfig: 1,
    EffectCircularConfig: 2,
    EffectRandomTwoColorInterpolationConfig: 3,
    EffectSwipeAndChange: 4,
    EffectReverseSwipe: 5,
}

mouse_led_effect_to_message_id = {
    EffectStaticConfig: 6,
    EffectCircularConfig: 7,
    EffectRandomTwoColorInterpolationConfig: 8,
    EffectSwipeAndChange: 9,
    EffectReverseSwipe: 10,
}

shelve_led_effect_to_message_id = {
    EffectStaticConfig: 15,
    EffectCircularConfig: 16,
    EffectRandomTwoColorInterpolationConfig: 17,
    EffectSwipeAndChange: 18,
    EffectReverseSwipe: 19,
    EffectStaticDetailedConfig: 20,
}

mouse_leds_index_ranges = {
    TouchButton.RIGHT_FOOT: (0, 6),
    TouchButton.LEFT_FOOT: (6, 6 + 6),
    TouchButton.LEFT_EAR: (6 + 6, 6 + 6 + 16),
    TouchButton.RIGHT_EAR: (6 + 6 + 16, 6 + 6 + 16 + 17),
}

PREV_BUTTON_LED_MSG = 21
NEXT_BUTTON_LED_MSG = 22


class RfidTokenRead:
    def __init__(self, id: bytes):
        self.id = id

    def __repr__(self):
        return "RFID Token (" + " ".join(f"{v:02x}" for v in self.id) + ")"


class RotaryEncoderEvent:
    def __init__(self, msg_content: bytes):
        self.position, self.increment, self.direction = struct.unpack("<iiB", msg_content)

    def __repr__(self):
        return f"Rotary event: pos {self.position}, incr {self.increment}, dir {self.direction}"


class TouchButtonPress:
    def __init__(self, msg_content: bytes):
        val = int(msg_content[0])
        self.touch_button = TouchButton(val)

    def __repr__(self) -> str:
        return "Pressed " + repr(self.touch_button)


class TouchButtonRelease:
    def __init__(self, msg_content: bytes):
        val = int(msg_content[0])
        self.touch_button = TouchButton(val)

    def __repr__(self) -> str:
        return "Released " + repr(self.touch_button)


class ButtonEvent:
    button_name = {1: 'left', 2: 'right', 3: 'rotary'}
    event_name = {
        0: 'pressed',
        1: 'released',
        2: 'clicked',
        3: 'double_clicked',
        4: 'long_pressed',
        5: 'repeat_pressed',
        6: 'long_released'
    }

    def __init__(self, msg_content: bytes):
        button_nr, event_nr = struct.unpack("<BB", msg_content)
        self.button = self.button_name[button_nr]
        self.event = self.event_name[event_nr]

    def __repr__(self) -> str:
        return f"Button {self.button} {self.event}"


incomingMsgMap = {
    0: RfidTokenRead,
    1: RotaryEncoderEvent,
    2: TouchButtonPress,
    3: TouchButtonRelease,
    4: ButtonEvent,
}


class MusicMouseProtocol(asyncio.Protocol):
    def __init__(self):
        super()
        self._msg_callback = None

    def register_message_callback(self, cb):
        self._msg_callback = cb

    def connection_made(self, transport):
        self.transport = transport
        self.in_buff = bytes()

    def __led_effect(self, effect_cfg, msg_dict):
        msg_content = effect_cfg.as_bytes()
        header = struct.pack("<IBH", MAGIC_TOKEN_HOST_TO_FW, msg_dict[type(effect_cfg)],
                             len(msg_content))
        self.transport.write(header + msg_content)

    def led_ring_effect(self, effect_cfg):
        self.__led_effect(effect_cfg, led_ring_effect_to_message_id)

    def mouse_led_effect(self, effect_cfg):
        self.__led_effect(effect_cfg, mouse_led_effect_to_message_id)

    def shelve_led_effect(self, effect_cfg):
        self.__led_effect(effect_cfg, shelve_led_effect_to_message_id)

    def button_background_led_prev(self, val):
        msg_content = struct.pack("<f", val)
        header = struct.pack("<IBH", MAGIC_TOKEN_HOST_TO_FW, PREV_BUTTON_LED_MSG, len(msg_content))
        self.transport.write(header + msg_content)

    def button_background_led_next(self, val):
        msg_content = struct.pack("<f", val)
        header = struct.pack("<IBH", MAGIC_TOKEN_HOST_TO_FW, NEXT_BUTTON_LED_MSG, len(msg_content))
        self.transport.write(header + msg_content)

    def data_received(self, data):
        self.in_buff += data
        self._parse_message()

    def connection_lost(self, exc):
        print('port closed')
        self.transport.loop.stop()

    def pause_writing(self):
        print('pause writing')
        print(self.transport.get_write_buffer_size())

    def resume_writing(self):
        print(self.transport.get_write_buffer_size())
        print('resume writing')

    def _parse_message(self):
        HEADER_SIZE = 4 + 1 + 2
        if len(self.in_buff) == 0:
            return
        if len(self.in_buff) >= HEADER_SIZE:
            token, msg_type, msg_size = struct.unpack("<IBH", self.in_buff[:HEADER_SIZE])
            if token == MAGIC_TOKEN_FW_TO_HOST and len(self.in_buff) >= HEADER_SIZE + msg_size:
                self._on_msg_receive(msg_type, self.in_buff[HEADER_SIZE:HEADER_SIZE + msg_size])
                self.in_buff = self.in_buff[HEADER_SIZE + msg_size:]
            else:
                idx = self.in_buff.find("\n".encode())
                if idx >= 0:
                    text_msg = self.in_buff[:idx]
                    print("LOG:", text_msg.decode())
                    self.in_buff = self.in_buff[idx + 1:]

    def _on_msg_receive(self, msg_type, msg_payload):
        parsed_msg = incomingMsgMap[msg_type](msg_payload)
        if self._msg_callback is not None:
            self._msg_callback(self, parsed_msg)