#!/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):
        self.session = session
        self.motion = session.service("ALMotion")
        self.posture = session.service("ALRobotPosture")
        self.video = None
        self.video_client = None
        self.running = True

        print("🔍 Initializing video...")
        self._init_video()

        pygame.init()
        self.screen = pygame.display.set_mode((640, 480))
        pygame.display.set_caption("NAO Teleoperation")
        self.clock = pygame.time.Clock()

    def _init_video(self):
        try:
            self.video = self.session.service("ALVideoDevice")
            print("✅ ALVideoDevice service found")
        except Exception as e:
            print(f"❌ Video service error: {e}")
            return

        # Try subscribeCamera (newer method)
        try:
            self.video_client = self.video.subscribeCamera("TeleopCam", 0, 2, 11, 10)  # camera 0 = top
            print("✅ Camera subscribed using subscribeCamera!")
            return
        except Exception as e:
            print(f"subscribeCamera failed: {e}")

        # Fallback to old subscribe
        try:
            self.video_client = self.video.subscribe("TeleopCam", 2, 11, 10)
            print("✅ Camera subscribed using old method")
        except Exception as e:
            print(f"❌ All subscription methods failed: {e}")

    def get_image(self):
        if not self.video or not self.video_client:
            return None
        try:
            image = self.video.getImageRemote(self.video_client)
            if image and len(image) >= 7:
                w, h = image[0], image[1]
                arr = np.frombuffer(bytearray(image[6]), dtype=np.uint8).reshape((h, w, 3))
                return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
        except:
            pass
        return None

    def run(self):
        self.motion.wakeUp()
        self.posture.goToPosture("StandInit", 0.5)
        print("🎮 Controls Active - WASD/Arrows | ESC to quit")

        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

            keys = pygame.key.get_pressed()
            x = y = theta = 0.0
            if keys[pygame.K_w] or keys[pygame.K_UP]:    x = 0.8
            if keys[pygame.K_s] or keys[pygame.K_DOWN]:  x = -0.8
            if keys[pygame.K_a] or keys[pygame.K_LEFT]:  theta = 0.6
            if keys[pygame.K_d] or keys[pygame.K_RIGHT]: theta = -0.6

            if abs(x) > 0.1 or abs(theta) > 0.1:
                self.motion.moveToward(x, y, theta)
            else:
                self.motion.stopMove()

            frame = self.get_image()
            if frame is not None:
                frame = cv2.resize(frame, (640, 480))
                rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                surf = pygame.surfarray.make_surface(rgb.swapaxes(0,1))
                self.screen.blit(surf, (0, 0))
            else:
                self.screen.fill((30, 30, 50))
                font = pygame.font.SysFont(None, 38)
                self.screen.blit(font.render("Camera Feed Not Available", True, (255, 100, 100)), (100, 200))

            pygame.display.flip()
            self.clock.tick(25)

        # Clean shutdown
        print("🛑 Shutting down...")
        self.motion.stopMove()
        self.motion.rest()
        if self.video and self.video_client:
            try:
                self.video.unsubscribe(self.video_client)
            except:
                pass
        pygame.quit()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--ip", type=str, default="127.0.0.1")
    parser.add_argument("--port", type=int, default=9559)
    args = parser.parse_args()

    session = qi.Session()
    try:
        session.connect(f"tcp://{args.ip}:{args.port}")
        print(f"✅ Connected to {args.ip}")
    except Exception as e:
        print(f"Connection failed: {e}")
        sys.exit(1)

    NaoTeleop(session).run()