# Import modules
from pybricks.hubs import PrimeHub
from pybricks.iodevices import PUPDevice
from pybricks.parameters import Button, Color, Direction, Port, Side, Stop
from pybricks.pupdevices import ColorSensor, ForceSensor, Motor, UltrasonicSensor
from pybricks.robotics import DriveBase
from pybricks.tools import multitask, run_task, StopWatch, wait
from umath import cos, radians
from urandom import randint

# Initialise program constants
SERVER = "http://192.168.0.129:5000"
DEVICE_NAME = "ROBOT1" # <--- Change device name in 3 char according to list below 

# Define channel map for BLE broadcast
CHANNEL_MAP = {
    "JKG": 1, # Jurassic Kingdom gate
    "TOI": 2, # Toilet
    "MOF": 3, # Moses & feeding
    "HAB": 4, # Hamster ball
    "VOL": 5, # Volcano
    "AAP": 6, # Aaron & Pedro
}

POLL_INTERVAL = 50 # Delay to poll server in ms
HEARTBEAT_INTERVAL = 2.0 # Delay to send heartbeat in s

# Initialise program variables
waiting = True

# ---------------
# BLE COMMS class
# ---------------
class BLEComms:
    # Initialise self
    def __init__(self, hub, device_name, broadcast_channel, observe_channels):
        print("[BLECOMMS] Initializing")
        self.hub = hub
        self.device_name = device_name
        self.broadcast_channel = broadcast_channel
        self.observe_channels = observe_channels
        
    # Send event
    def broadcast(self, kind, value):
        # Populate message
        msg = self.device_name + ":" + kind + ":" + value

        payload = msg[:21] # BLE advertisement limit of 21 char

        print("[BLE] Sending payload '", payload, "'")
        # Send 3 times for reliability
        for i in range(3):
            self.hub.ble.broadcast(payload)
            wait(POLL_INTERVAL)
    
    # Receiving event    
    def observe(self):
        for ch in self.observe_channels:
            msg = self.hub.ble.observe(ch)
            
            if msg:
                print("[BLE] Receiving message '", msg, "'")
                return msg
        
        return None

# --------------
# Hardware class
# --------------
class Hardware:
    # Initialise self    
    def __init__(self): # <--- Add motors and sensors as required
        print("[HARDWARE] Initializing")
        # Motors
        self.motor_a = Motor(Port.A) 
        self.motor_b = Motor(Port.B)
        # Sensors
        self.sensor_us = UltrasonicSensor(Port.E)

    # Example action
    def do_action(self): # <--- Add motors and sensors as required
        self.motor_a.run_angle(500, 90)
        self.motor_b.run_angle(500, -90)
        
    def stop(self): # <--- Add motors and sensors as required
        self.motor_a.stop()
        self.motor_b.stop()

# ------------
# Device class
# ------------
class Device:
    # Initialise self
    def __init__(self, name):
        self.device_name = name
        broadcast_channel = CHANNEL_MAP[name]

        # Initialise array for all channels except own
        observe_channels = [
            ch for dev, ch in CHANNEL_MAP.items()
            if dev != name
        ]
        self.hub = PrimeHub(
            broadcast_channel=broadcast_channel,
            observe_channels=observe_channels # BLE observe on all channels except own
        )
        self.ble = BLEComms(
            self.hub, 
            name, 
            broadcast_channel, 
            observe_channels
        )
        self.hw = Hardware()
        self.clock = StopWatch()
        self.last_heartbeat = 0
        
    # Send status 'alive' to Flask server
    def heartbeat(self):
        self.ble.broadcast("STATUS", "alive")

    # Define functions run periodically
    def periodic(self): # <--- Add required code        
        pass
        
    # Handle BLE message handler
    def on_ble_message(self, msg):
        try:
            src, kind, value = msg.split(":", 2)
        except:
            return

        # Ignore empty or own broadcasts
        if not src or src == self.device_name:
            return

        print(f"[EVENT] Receiving message from device {src}: {kind}:{value}")

        if "_" in value:
            device, action = value.split("_", 1)
        else:
            device = value
            action = None
        
        if device == self.device_name and kind == "CMD": # <--- Add required commands & return messages
            # Example: Gate tells Fossil Dig to start
            if action == "start":
                self.hw.do_action()
                self.ble.broadcast("EVENT", "digging")
            # Stop command
            if action == "stop":
                self.hw.stop()
                self.ble.broadcast("EVENT", "stopping")
       
    # Send event 'ready' to Flask server
    def startup(self):
        print("[DEVICE] Starting up")
        self.ble.broadcast("EVENT", "ready")
        
    # Define main function
    def run(self):
        # Send event        
        self.startup()
        
        while True:
            # Send status
            now = self.clock.time()
            if now - self.last_heartbeat > HEARTBEAT_INTERVAL:
                self.heartbeat()
                self.last_heartbeat = now
                                            
            # Listen for BLE broadcast
            msg = self.ble.observe()
            if msg:
                self.on_ble_message(msg)
                               
            self.periodic()

            wait(POLL_INTERVAL)

# Declare main program
if __name__ == "__main__":
    
    device = Device(
        name=DEVICE_NAME
    )
    
    device.run()