from shop import Shop
from wand import Wand
from constants import *
import asyncio
import sys

# Custom wand class extending the default wand
class MyWand(Wand):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.colors = ["#a333c8", "2185d0", "0x21ba45", "#fbbd08", "#f2711c", "#db2828"]
        self.position_id = None

    # Do some functions after connecting
    def post_connect(self):
        print("Connected to {}".format(self.name))
        # # Vibrate the wand and set its color to red
        # self.set_led(self.colors.pop())
        # # Subscribe to notifications
        self.subscribe_button()

    # Button callback, automatically called after connecting to wand
    async def on_button(self, pressed):
        if pressed:
            # Unsubscribe from the position callback
            await self.set_led(self.colors.pop())
            # Disconnect if we run out of colors
            if len(self.colors) == 0:
                await self.disconnect()

async def main():

    # Create a new wand scanner
    shop = Shop(asyncio.get_running_loop(), wand_class=MyWand, debug=True)
    wands = []
    try:
        # While we don't have any wands
        while len(wands) == 0:
            # Scan for wands and automatically connect
            print("Scanning...")
            wands = await shop.scan(connect=True)
            # For each wand (Only tested with one)
            for wand in wands:
                # Vibrate the wand and set its color to red
                await wand.vibrate(PATTERN.BURST)

                # Callback for position
                # def onPos(x, y, pitch, roll):
                #     pitch = "Pitch: {}".format(pitch).ljust(16)
                #     roll = "Roll: {}".format(roll).ljust(16)
                #     print("{}{}(x, y): ({}, {})".format(pitch, roll, x, y))

                # Add the event callback to the wand
                # wand.position_id = wand.on("position", onPos)

    # Detect keyboard interrupt disconnect
    except KeyboardInterrupt as e:
        for wand in wands:
            wand.disconnect()

if __name__ == "__main__":
    asyncio.run(main())