# 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 = "AAP" # <--- 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

# Open, lift, rotate & flap angle, motor speed in deg/s, time to wait in ms, list of motors
ANGLE_OPEN = 120
ANGLE_LIFT = 120
ANGLE_ROTATE = 60
ANGLE_FLAP = 360
SPEED_MOTOR = 360
WAIT_TIME = 100

# ---------------
# 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_aaron = Motor(Port.A, Direction.CLOCKWISE) # Aaron's head
        self.motor_pedro_bottom = Motor(Port.B, Direction.CLOCKWISE) # Pedro's lift
        #self.motor_moses_top = Motor(Port.C, Direction.CLOCKWISE) # Moses' mouth
        self.motor_pedro_top = Motor(Port.D, Direction.CLOCKWISE) # Pedro's wings
        # Sensors
        self.sensor_us_back = UltrasonicSensor(Port.E) # Backup ultra-sonic sensor facing back
        self.sensor_us = UltrasonicSensor(Port.F) # Main ultra-sonic sensor facing front
        # Motor list for reset() method
        self.list_motors = [self.motor_aaron,self.motor_pedro_bottom,'',self.motor_pedro_top,'',''] # Port A-F

    # Example action
    def do_action(self): # <--- Add motors and sensors as required
        self.motor_aaron.run_angle(500, 90)
        self.motor_pedro_bottom.run_angle(500, -90)
        
    def stop(self): # <--- Add motors and sensors as required
        self.motor_aaron.stop()
        self.motor_pedro_bottom.stop()
        self.motor_pedro_top.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
        
        self.counter_us = 0
        self.frequency = "low"
        self.aaron_attack_ready = True
        self.frequency = "low"
        self.pedro_action = None

    # Make Aaron jump out of cave, bite Indie & go back in cave
    def bite(self):
        print("[AARON] Bite triggered!")

        m = self.hw.motor_aaron

        # Jump out
        m.run_angle(SPEED_MOTOR, ANGLE_OPEN)

        wait(WAIT_TIME * 5)

        # Bite motion
        m.run_angle(SPEED_MOTOR, -ANGLE_ROTATE)
        m.run_angle(SPEED_MOTOR, ANGLE_ROTATE)

        wait(WAIT_TIME * 5)

        # Return to cave
        m.run_angle(SPEED_MOTOR, -ANGLE_OPEN)

        # Broadcast event
        self.ble.broadcast("EVENT", "AAP_bite_done")
    
    # Make Pedro flap wings, lift up & down
    def flap_lift(self, mode):
        """
        mode = "flap", "lift", or "flap_lift"
        """

        speed = SPEED_MOTOR
        angle_flap = ANGLE_FLAP
        angle_lift = ANGLE_LIFT

        pedro_top = self.hw.motor_pedro_top
        pedro_bottom = self.hw.motor_pedro_bottom

        print("[PEDRO] Performing", mode)

        # Number of repetitions
        if mode == "flap":
            reps = 6
        elif mode == "lift":
            reps = 2
        elif mode == "flap_lift":
            reps = 2
        else:
            return

        # Initial movement
        if mode == "flap":
            pedro_top.run_target(speed, angle_flap)
        elif mode == "lift":
            pedro_bottom.run_target(speed, angle_lift)
        elif mode == "flap_lift":
            pedro_top.run_target(speed, angle_flap, wait=False)
            pedro_bottom.run_target(speed, angle_lift)

        # Repeated movement
        for _ in range(reps):
            if mode == "flap":
                pedro_top.run_angle(speed, angle_flap)
            elif mode == "lift":
                pedro_bottom.run_angle(speed, angle_lift)
            elif mode == "flap_lift":
                pedro_top.run_angle(speed, angle_flap, wait=False)
                pedro_bottom.run_angle(speed, angle_lift)

            # Toggle direction
            angle_flap = -angle_flap
            angle_lift = -angle_lift

        # Reset to neutral
        if mode == "flap":
            pedro_top.run_target(speed, 0)
        elif mode == "lift":
            pedro_bottom.run_target(speed, 0)
        elif mode == "flap_lift":
            pedro_top.run_target(speed, 0, wait=False)
            pedro_bottom.run_target(speed, 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        
        
        # Pedro BLE-triggered action
        if self.pedro_action == "flap":
            self.flap_lift("flap")
        elif self.pedro_action == "lift":
            self.flap_lift("lift")
        elif self.pedro_action == "flap_lift":
            self.flap_lift("flap_lift")
            self.pedro_action = None

        # Aaron ultrasonic-triggered attack
        dist = self.hw.sensor_us.distance()
        if dist is not None and dist < 200 and self.aaron_attack_ready:
            self.bite()
            self.aaron_attack_ready = False

        # High-frequency mode
        if self.frequency == "high":
            # optional: extra flapping or shaking
            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
            if action == "start":
                #self.open_close_gate("open")
                self.ble.broadcast("EVENT", "start")
            if action == "flap":
                self.pedro_action = "flap"
                self.ble.broadcast("EVENT", "flap")
            elif action == "lift":
                self.pedro_action = "lift"
                self.ble.broadcast("EVENT", "lift")
            elif action == "flap_lift":
                self.pedro_action = "flap_lift"
                self.ble.broadcast("EVENT", "flap_lift")
            elif action == "fast":
                self.frequency = "high"
                self.ble.broadcast("EVENT", "high")
            elif action == "stop":
                motors_stop = [0,0,"",0,"",""] # Port A-F
                self.reset(motors_stop)
                self.ble.broadcast("EVENT", "stopping")
                
            # Example: Gate tells Fossil Dig to start
            #if action == "start":
            #    self.hw.do_action()
            #    self.ble.broadcast("EVENT", "digging")
            # Stop command            

    # Reset Aaron & Pedro motors
    def reset(self, motors):
        # Initialize counter
        counter_motor = 0
        # Loop through motors array
        for angle in motors:
            # If a value is found in array
            if angle != '':
                print('Resetting motor', self.hw.list_motors[counter_motor], 'at', angle, 'degrees ...')
                # Look up motor, reset position to angle
                self.hw.list_motors[counter_motor].run_target(SPEED_MOTOR, angle)
            # Increment counter
            counter_motor += 1
               
    # Send event 'ready' to Flask server
    def startup(self):
        print("[DEVICE] Starting up")
        self.ble.broadcast("EVENT", "ready")
        
    # Define main function
    def run(self):
        # Reset Aaron, Pedro motors, set angle for all active motors
        motors_start = [0,-90,'',0,'',''] # Port A-F        
        # Send event        
        self.startup()
        self.reset(motors_start)
        
        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()