from codrone_edu.drone import *
import pygame
import time

# ====================== SETTINGS ======================
POWER = 22          # Safe & slow speed (15-30)
DEADZONE = 0.15
# =====================================================

pygame.init()
screen = pygame.display.set_mode((500, 300))
pygame.display.set_caption("CoDrone Controller")

# Gamepad setup
pygame.joystick.init()
joysticks = [pygame.joystick.Joystick(i) for i in range(pygame.joystick.get_count())]
for joy in joysticks:
    joy.init()
    print(f"🎮 Detected: {joy.get_name()}")

drone = Drone()
drone.connect()
time.sleep(3)

try:
    drone.takeoff()
    print("✅ Drone flying! Controls: Arrows / Left Stick | W/S | A/D")
    print("ESC or Start button to land")

    clock = pygame.time.Clock()

    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT or \
               (event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE) or \
               (event.type == pygame.JOYBUTTONDOWN and event.button in (7, 6)):  # Start/Options
                running = False

        keys = pygame.key.get_pressed()
        roll = pitch = throttle = yaw = 0

        # ==================== KEYBOARD ====================
        if keys[pygame.K_UP]:    pitch = POWER
        if keys[pygame.K_DOWN]:  pitch = -POWER
        if keys[pygame.K_LEFT]:  roll = -POWER
        if keys[pygame.K_RIGHT]: roll = POWER
        if keys[pygame.K_w]:     throttle = POWER
        if keys[pygame.K_s]:     throttle = -POWER
        if keys[pygame.K_a]:     yaw = -POWER
        if keys[pygame.K_d]:     yaw = POWER

        # ==================== GAMEPAD ====================
        if joysticks:
            joy = joysticks[0]

            # Left stick
            roll += int(joy.get_axis(0) * POWER)
            pitch += int(-joy.get_axis(1) * POWER)

            # Right stick / Triggers
            if joy.get_numaxes() > 3:
                throttle += int(-joy.get_axis(3) * POWER)
                yaw += int(joy.get_axis(2) * POWER)

            # Common trigger mapping (Xbox/PS)
            if joy.get_numaxes() > 5:
                yaw += int(joy.get_axis(5) * POWER)   # Right trigger
                yaw += int(-joy.get_axis(4) * POWER)  # Left trigger

        # Clamp values
        roll = max(-100, min(100, roll))
        pitch = max(-100, min(100, pitch))
        throttle = max(-100, min(100, throttle))
        yaw = max(-100, min(100, yaw))

        # Apply movement
        drone.set_roll(roll)
        drone.set_pitch(pitch)
        drone.set_throttle(throttle)
        drone.set_yaw(yaw)
        drone.move(0.1)

        clock.tick(30)

except Exception as e:
    print("Error:", e)

finally:
    print("Landing...")
    drone.land()
    drone.disconnect()
    pygame.quit()
    print("Disconnected.")