#!/usr/bin/env python
# -*- encoding: UTF-8 -*-

import qi
import argparse
import sys
import time
import pygame
import cv2
import numpy as np

class NaoTeleop:
    def __init__(self, session, robot_ip="127.0.0.1", port=9559):
        self.motion = session.service("ALMotion")
        self.posture = session.service("ALRobotPosture")
        self.video = session.service("ALVideoDevice")
        
        self.robot_ip = robot_ip
        self.running = True
        
        # Video settings (using numbers instead of vision_definitions)
        self.resolution = 2   # 2 = kQVGA (320x240)
        self.color_space = 11 # 11 = kRGBColorSpace
        self.fps = 15
        
        self.video_client = None
        try:
            self.video_client = self.video.subscribe("TeleopCamera", self.resolution, self.color_space, self.fps)
            print("✅ Camera subscribed successfully")
        except Exception as e:
            print(f"⚠️ Camera subscription failed: {e}")
        
        # Pygame setup
        pygame.init()
        self.screen = pygame.display.set_mode((640, 480))
        pygame.display.set_caption("NAO Teleoperation - Camera + Controls")
        self.clock = pygame.time.Clock()
        
        # Movement
        self.speed = 0.8
        self.turn_speed = 0.6
        self.strafe = 0.5

    def get_image(self):
        if not self.video_client:
            return None
        try:
            image = self.video.getImageRemote(self.video_client)
            if not image:
                return None
            
            width = image[0]
            height = image[1]
            array = image[6]
            
            image_array = np.frombuffer(bytearray(array), dtype=np.uint8).reshape((height, width, 3))
            return cv2.cvtColor(image_array, cv2.COLOR_RGB2BGR)
        except:
            return None

    def handle_input(self):
        keys = pygame.key.get_pressed()
        x = y = theta = 0.0

        # Keyboard
        if keys[pygame.K_w] or keys[pygame.K_UP]:    x = self.speed
        if keys[pygame.K_s] or keys[pygame.K_DOWN]:  x = -self.speed
        if keys[pygame.K_a] or keys[pygame.K_LEFT]:  theta = self.turn_speed
        if keys[pygame.K_d] or keys[pygame.K_RIGHT]: theta = -self.turn_speed
        if keys[pygame.K_q]: y = self.strafe
        if keys[pygame.K_e]: y = -self.strafe

        # Game Controller
        if pygame.joystick.get_count() > 0:
            joy = pygame.joystick.Joystick(0)
            joy.init()
            
            x = -joy.get_axis(1) * self.speed      # Left stick Y
            y = joy.get_axis(0) * self.strafe      # Left stick X
            theta = -joy.get_axis(3) * self.turn_speed  # Right stick X (turning)

            if joy.get_button(0):  # A / X button
                self.posture.goToPosture("StandInit", 0.5)

        # Send command
        if abs(x) > 0.1 or abs(y) > 0.1 or abs(theta) > 0.1:
            self.motion.moveToward(x, y, theta, [["Frequency", 1.0]])
        else:
            self.motion.stopMove()

    def run(self):
        self.motion.wakeUp()
        self.posture.goToPosture("StandInit", 0.5)
        self.motion.setMoveArmsEnabled(True, True)
        
        print("\n🎮 NAO Teleop Started!")
        print("WASD / Arrows = Move & Turn")
        print("Q/E = Strafe left/right")
        print("ESC = Quit")
        
        try:
            while self.running:
                for event in pygame.event.get():
                    if event.type == pygame.QUIT or (event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE):
                        self.running = False
                
                self.handle_input()
                
                frame = self.get_image()
                if frame is not None:
                    frame = cv2.resize(frame, (640, 480))
                    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    surf = pygame.surfarray.make_surface(frame_rgb.swapaxes(0, 1))
                    self.screen.blit(surf, (0, 0))
                else:
                    # Show placeholder if no camera
                    self.screen.fill((30, 30, 30))
                    font = pygame.font.SysFont(None, 48)
                    text = font.render("Waiting for NAO Camera...", True, (255, 100, 100))
                    self.screen.blit(text, (100, 200))
                
                pygame.display.flip()
                self.clock.tick(30)
        
        finally:
            print("🛑 Shutting down...")
            self.motion.stopMove()
            self.motion.rest()
            if self.video_client:
                self.video.unsubscribe(self.video_client)
            pygame.quit()

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--ip", type=str, default="127.0.0.1", help="Robot IP address")
    parser.add_argument("--port", type=int, default=9559, help="Naoqi port")
    
    args = parser.parse_args()
    
    session = qi.Session()
    try:
        session.connect(f"tcp://{args.ip}:{args.port}")
        print(f"✅ Connected to NAO at {args.ip}:{args.port}")
    except RuntimeError:
        print(f"❌ Can't connect to Naoqi at {args.ip}:{args.port}")
        sys.exit(1)
    
    teleop = NaoTeleop(session, args.ip, args.port)
    teleop.run()

if __name__ == "__main__":
    main()