from codrone_edu.drone import *
import pygame
import time
import sys

# Initialize pygame
pygame.init()
screen = pygame.display.set_mode((400, 300))  # Small window (focus here)
pygame.display.set_caption("CoDrone Arrow Key Controller")

drone = Drone()
drone.connect()   # or drone.pair()

print("Connecting...")
time.sleep(3)

try:
    drone.takeoff()
    print("✅ Drone ready!")
    print("Use Arrow Keys:")
    print("↑ Forward  ↓ Backward  ← Left  → Right")
    print("W = Up    S = Down")
    print("Press ESC or close window to land")

    power = 20          # Very safe & slow (15-25 recommended)

    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):
                running = False

        # Get current key states
        keys = pygame.key.get_pressed()

        # Reset all movement
        drone.set_roll(0)
        drone.set_pitch(0)
        drone.set_yaw(0)
        drone.set_throttle(0)

        # Arrow keys for horizontal movement
        if keys[pygame.K_UP]:
            drone.set_pitch(power)      # Forward
        if keys[pygame.K_DOWN]:
            drone.set_pitch(-power)     # Backward
        if keys[pygame.K_LEFT]:
            drone.set_roll(-power)      # Left
        if keys[pygame.K_RIGHT]:
            drone.set_roll(power)       # Right

        # W/S for altitude
        if keys[pygame.K_w]:
            drone.set_throttle(power)   # Up
        if keys[pygame.K_s]:
            drone.set_throttle(-power)  # Down

        drone.move(0.1)   # Send command
        clock.tick(30)    # ~30 updates per second (smooth but not too fast)

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

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