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

import qi
import argparse
import sys
import time
import threading
import queue
import cv2
import numpy as np
import pygame
import yt_dlp
import os

# --- Audio for mic ---
try:
    import pyaudio
    HAS_AUDIO = True
except ImportError:
    HAS_AUDIO = False

# ==========================================
# NAO MIC STREAM
# ==========================================
class SoundReceiver:
    PREBUFFER_CHUNKS = 3
    SILENCE_CHUNK = b"\x00" * 4096

    def __init__(self, output_device_index=None, mic_gain=8.0):
        self.running = True
        self.underrun_count = 0
        self.chunks_received = 0
        self._write_errors = 0
        self._peak_since_print = 0
        self._last_status_print = 0
        self.mic_gain = float(mic_gain)

        if HAS_AUDIO:
            self.p = pyaudio.PyAudio()
            print("🔈 Available output devices:")
            for i in range(self.p.get_device_count()):
                info = self.p.get_device_info_by_index(i)
                if info.get("maxOutputChannels", 0) > 0:
                    print(f"   [{i}] {info['name']}")
            try:
                chosen = (self.p.get_device_info_by_index(output_device_index) if output_device_index is not None
                          else self.p.get_default_output_device_info())
                print(f"🔈 Using: [{chosen['index']}] {chosen['name']}")
            except: pass

            self.stream = self.p.open(format=pyaudio.paInt16, channels=1, rate=16000,
                                      output=True, output_device_index=output_device_index,
                                      frames_per_buffer=2048)
            self._play_test_tone()
            self.queue = queue.Queue(maxsize=40)
            self.playback_thread = threading.Thread(target=self._playback_loop, daemon=True)
            self.playback_thread.start()

    def _play_test_tone(self):
        try:
            duration, freq, rate = 0.3, 440.0, 16000
            t = np.linspace(0, duration, int(rate * duration), False)
            tone = (np.sin(freq * t * 2 * np.pi) * 12000).astype(np.int16)
            self.stream.write(tone.tobytes())
        except: pass

    def processRemote(self, nbOfChannels, nbOfSamplesByChannel, timeStamp, buffer):
        if not HAS_AUDIO: return
        self.chunks_received += 1
        try:
            self.queue.put_nowait(bytes(buffer))
        except queue.Full:
            try:
                self.queue.get_nowait()
                self.queue.put_nowait(bytes(buffer))
            except: pass

    def _playback_loop(self):
        primed = []
        while self.running and len(primed) < self.PREBUFFER_CHUNKS:
            try:
                primed.append(self.queue.get(timeout=1.0))
            except queue.Empty:
                break
        for chunk in primed:
            self._write_chunk(chunk)

        while self.running:
            try:
                chunk = self.queue.get(timeout=0.2)
                self._write_chunk(chunk)
            except queue.Empty:
                self.underrun_count += 1
                try:
                    self.stream.write(self.SILENCE_CHUNK)
                except: pass

            now = time.time()
            if now - self._last_status_print > 10:
                self._last_status_print = now
                print(f"🎤 mic: {self.chunks_received} chunks, {self.underrun_count} underruns, peak {self._peak_since_print}")
                self._peak_since_print = 0

    def _write_chunk(self, raw_bytes):
        try:
            samples = np.frombuffer(raw_bytes, dtype=np.int16).astype(np.float32)
            if samples.size:
                amplified = np.clip(samples * self.mic_gain, -32768, 32767)
                peak = int(np.abs(amplified).max())
                if peak > self._peak_since_print:
                    self._peak_since_print = peak
                self.stream.write(amplified.astype(np.int16).tobytes())
            else:
                self.stream.write(raw_bytes)
        except: pass

    def close(self):
        self.running = False
        if HAS_AUDIO:
            try: self.playback_thread.join(timeout=1.0)
            except: pass
            try:
                self.stream.stop_stream()
                self.stream.close()
                self.p.terminate()
            except: pass


# ==========================================
# MUSIC PLAYER ON NAO
# ==========================================
class NaoMusicPlayer:
    def __init__(self, session):
        self.session = session
        self.player = session.service("ALAudioPlayer")
        self.music_dir = "/home/nao/music"
        try:
            self.session.service("ALFileManager").createFolder(self.music_dir)
        except: pass
        self.current_file = None
        self.is_playing = False
        self.current_title = ""

    def search(self, query):
        try:
            ydl_opts = {
                'format': 'bestaudio/best',
                'quiet': True,
                'noplaylist': True,
                'extract_flat': True,
            }
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                info = ydl.extract_info(f"ytsearch5:{query}", download=False)
                return info.get('entries', [])
        except:
            return []

    def play(self, url, title):
        self.stop()
        self.current_title = title
        try:
            ydl_opts = {
                'format': 'bestaudio/best',
                'outtmpl': '/tmp/%(id)s.%(ext)s',
                'quiet': True,
                'noplaylist': True,
                'extractaudio': True,
                'postprocessors': [{'key': 'FFmpegExtractAudio', 'preferredcodec': 'mp3'}],
                'ignoreerrors': True,
                'no_warnings': True,
            }
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                info = ydl.extract_info(url, download=True)
                if not info:
                    return False
                downloaded = ydl.prepare_filename(info)

            remote_path = f"{self.music_dir}/{int(time.time())}.mp3"
            file_manager = self.session.service("ALFileManager")
            # Use uploadFile instead of put
            file_manager.uploadFile(downloaded, remote_path)

            self.current_file = remote_path
            self.player.playFile(remote_path)
            self.is_playing = True
            print(f"♪ Playing on NAO: {title}")
            return True
        except Exception as e:
            print(f"Playback error: {e}")
            return False

    def pause(self):
        if not self.current_file: return
        if self.is_playing:
            self.player.pause()
            self.is_playing = False
        else:
            self.player.play(self.current_file)
            self.is_playing = True

    def stop(self):
        try:
            self.player.stopAll()
        except:
            pass
        self.is_playing = False
        self.current_file = None

    def seek(self, seconds):
        try:
            pos = self.player.getCurrentPosition() / 1000.0
            self.player.goTo(pos + seconds)
        except: pass


# ==========================================
# MAIN TELEOP
# ==========================================
class NaoTeleop:
    def __init__(self, session, audio_device_index=None, mic_channel=1, mic_gain=8.0, video_scale=2.0):
        self.session = session
        self.motion = session.service("ALMotion")
        self.posture = session.service("ALRobotPosture")
        self.tts = session.service("ALTextToSpeech")
        try:
            self.battery = session.service("ALBattery")
        except:
            self.battery = None

        self.music = NaoMusicPlayer(session)
        self.music_results = []
        self.selected_result = 0

        self.audio_device = None
        self.audio_service_name = "SoundReceiver"
        self.audio_device_index = audio_device_index
        self.mic_channel = mic_channel
        self.mic_gain = mic_gain

        self.running = True
        self.head_yaw = 0.0
        self.head_pitch = 0.0
        self.battery_level = 100
        self.last_batt_check = 0
        self.volume = 50

        self.typing_mode = False
        self.chat_message = ""
        self.music_search_mode = False
        self.music_search_text = ""

        self.video_scale = video_scale
        self.display_width = int(320 * video_scale)
        self.display_height = int(240 * video_scale)

        self._init_audio_device()
        self._init_video_maxfps()
        if HAS_AUDIO:
            self._init_mic_stream()

        pygame.init()
        self.screen = pygame.display.set_mode((self.display_width, self.display_height))
        pygame.display.set_caption("NAO Teleop + Music on Robot")
        self.font = pygame.font.SysFont(None, 24)
        self.clock = pygame.time.Clock()
        self.update_title()

    def _init_audio_device(self):
        try:
            self.audio_device = self.session.service("ALAudioDevice")
            self.volume = self.audio_device.getOutputVolume()
            print(f"🔊 NAO speaker volume: {self.volume}%")
        except Exception as e:
            print(f"⚠️ ALAudioDevice error: {e}")
            self.audio_device = None

    def _init_video_maxfps(self):
        try:
            self.video = self.session.service("ALVideoDevice")
            self.video_client = self.video.subscribeCamera("TeleopCam", 0, 0, 11, 25)
            print("✅ Camera ready")
        except:
            print("❌ Camera not available")
            self.video_client = None

    def _init_mic_stream(self):
        if not self.audio_device: return
        try:
            self.sound_receiver = SoundReceiver(self.audio_device_index, self.mic_gain)
            self.session.registerService(self.audio_service_name, self.sound_receiver)
            time.sleep(0.5)
            self.audio_device.setClientPreferences(self.audio_service_name, 16000, self.mic_channel, 0)
            self.audio_device.subscribe(self.audio_service_name)
            print(f"✅ Mic stream ready (gain={self.mic_gain}x)")
        except Exception as e:
            print(f"❌ Mic init failed: {e}")

    def get_image(self):
        if not hasattr(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))
                frame = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
                return cv2.resize(frame, (320, 240), interpolation=cv2.INTER_NEAREST)
        except:
            pass
        return None

    def wave_gesture(self):
        if self.typing_mode: return
        print("🤚 Waving hello...")
        try:
            self.motion.setStiffnesses("RArm", 1.0)
            names = ["RShoulderPitch", "RShoulderRoll", "RElbowYaw", "RElbowRoll", "RWristYaw", "RHand"]
            angles = [[-1.5]*5, [-0.4,-0.7,-0.4,-0.7,-0.4], [0.6,1.1,0.6,1.1,0.6],
                      [0.8,0.3,0.8,0.3,0.8], [0.0,1.5,0.0,-1.5,0.0], [1.0]*5]
            times = [[0.5,1.0,1.5,2.0,2.5]] * 6
            self.motion.angleInterpolation(names, angles, times, True)
            neutral = [1.5, -0.15, 1.2, 0.5, 0.0, 0.0]
            self.motion.angleInterpolationWithSpeed(names, neutral, 0.3)
            self.motion.setStiffnesses("RArm", 0.6)
        except Exception as e:
            print(f"Wave error: {e}")

    def handle_head_movement(self):
        if self.typing_mode or self.music_search_mode: return
        keys = pygame.key.get_pressed()
        speed = 0.3
        changed = False
        if keys[pygame.K_i]: self.head_pitch = max(self.head_pitch - speed, -0.5); changed = True
        if keys[pygame.K_k]: self.head_pitch = min(self.head_pitch + speed, 0.5); changed = True
        if keys[pygame.K_j]: self.head_yaw = min(self.head_yaw + speed, 2.0); changed = True
        if keys[pygame.K_l]: self.head_yaw = max(self.head_yaw - speed, -2.0); changed = True
        if changed:
            self.motion.setAngles(["HeadYaw", "HeadPitch"], [self.head_yaw, self.head_pitch], 0.3)

    def reset_head(self):
        self.head_yaw = 0.0
        self.head_pitch = 0.0
        self.motion.setAngles(["HeadYaw", "HeadPitch"], [0.0, 0.0], 0.4)

    def update_title(self):
        mode = "[TYPING] " if self.typing_mode else ""
        pygame.display.set_caption(f"NAO Teleop | {mode}Battery: {self.battery_level}% | Vol: {self.volume}%")

    def change_volume(self, delta):
        self.volume = int(min(100, max(0, self.volume + delta)))
        if self.audio_device:
            try:
                self.audio_device.setOutputVolume(self.volume)
            except: pass
        print(f"🔊 NAO volume: {self.volume}%")
        self.update_title()

    def check_battery(self):
        if not self.battery: return
        now = time.time()
        if now - self.last_batt_check > 10:
            try:
                self.battery_level = self.battery.getBatteryCharge()
            except: pass
            self.update_title()
            self.last_batt_check = now

    def run(self):
        self.motion.wakeUp()
        self.posture.goToPosture("StandInit", 0.5)
        time.sleep(1.0)
        self.motion.setMoveArmsEnabled(True, True)

        carpet_config = [["StepHeight", 0.028], ["MaxStepFrequency", 0.6], ["TorsoWy", 0.05]]

        print("\n🎮 Controls: WASD=Walk | IJKL=Head | SPACE=Reset head | 1=Wave | M=Music search | P=Pause | X=Stop | ←→=Seek | T=TTS | ESC=Quit")

        while self.running:
            self.check_battery()

            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    self.running = False
                elif event.type == pygame.KEYDOWN:
                    if self.music_search_mode:
                        if event.key == pygame.K_RETURN:
                            if self.music_search_text.strip():
                                self.music_results = self.music.search(self.music_search_text)
                                if self.music_results:
                                    self.selected_result = 0
                                else:
                                    self.music_search_mode = False
                            else:
                                self.music_search_mode = False
                        elif event.key == pygame.K_ESCAPE:
                            self.music_search_mode = False
                            self.music_search_text = ""
                        elif event.key == pygame.K_BACKSPACE:
                            self.music_search_text = self.music_search_text[:-1]
                        elif event.key in (pygame.K_1, pygame.K_2, pygame.K_3, pygame.K_4, pygame.K_5):
                            idx = int(event.unicode) - 1
                            if idx < len(self.music_results):
                                entry = self.music_results[idx]
                                self.music.play(entry['url'], entry.get('title', 'Unknown'))
                                self.music_search_mode = False
                        else:
                            self.music_search_text += event.unicode
                    elif self.typing_mode:
                        if event.key == pygame.K_RETURN:
                            if self.chat_message.strip():
                                threading.Thread(target=self.tts.say, args=(self.chat_message,)).start()
                            self.chat_message = ""
                            self.typing_mode = False
                        elif event.key == pygame.K_ESCAPE:
                            self.chat_message = ""
                            self.typing_mode = False
                        elif event.key == pygame.K_BACKSPACE:
                            self.chat_message = self.chat_message[:-1]
                        else:
                            self.chat_message += event.unicode
                    else:
                        if event.key == pygame.K_ESCAPE:
                            self.running = False
                        elif event.key == pygame.K_1:
                            self.wave_gesture()
                        elif event.key == pygame.K_SPACE:
                            self.reset_head()
                        elif event.key == pygame.K_m:
                            self.music_search_mode = True
                            self.music_search_text = ""
                            self.music_results = []
                        elif event.key == pygame.K_p:
                            self.music.pause()
                        elif event.key == pygame.K_x:   # Stop music
                            self.music.stop()
                        elif event.key == pygame.K_LEFT:
                            self.music.seek(-10)
                        elif event.key == pygame.K_RIGHT:
                            self.music.seek(10)
                        elif event.key == pygame.K_t:
                            self.typing_mode = True
                            self.motion.stopMove()

            if not self.typing_mode and not self.music_search_mode:
                keys = pygame.key.get_pressed()
                x = y = theta = 0.0
                if keys[pygame.K_w] or keys[pygame.K_UP]: x = 0.6
                if keys[pygame.K_s] or keys[pygame.K_DOWN]: x = -0.6
                if keys[pygame.K_a] or keys[pygame.K_LEFT]: theta = 0.5
                if keys[pygame.K_d] or keys[pygame.K_RIGHT]: theta = -0.5

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

            frame = self.get_image()
            if frame is not None:
                up = cv2.resize(frame, (self.display_width, self.display_height), interpolation=cv2.INTER_CUBIC)
                rgb = cv2.cvtColor(up, cv2.COLOR_BGR2RGB)
                surf = pygame.surfarray.make_surface(rgb.swapaxes(0,1))
                self.screen.blit(surf, (0, 0))
            else:
                self.screen.fill((10, 10, 30))

            # Music status
            if self.music.current_title:
                status = f"♪ {self.music.current_title[:45]} {'(Paused)' if not self.music.is_playing else ''}"
                text = self.font.render(status, True, (0, 255, 100))
                self.screen.blit(text, (10, 10))

            # Music search interface
            if self.music_search_mode:
                s = pygame.Surface((self.display_width, self.display_height//2))
                s.set_alpha(220)
                s.fill((0, 0, 40))
                self.screen.blit(s, (0, 80))

                text = self.font.render(f"Search: {self.music_search_text}", True, (255, 255, 255))
                self.screen.blit(text, (20, 100))

                if self.music_results:
                    for i, entry in enumerate(self.music_results[:5]):
                        color = (255, 255, 0) if i == self.selected_result else (200, 200, 200)
                        title = entry.get('title', 'No title')[:60]
                        line = self.font.render(f"{i+1}. {title}", True, color)
                        self.screen.blit(line, (20, 140 + i*30))
                else:
                    text = self.font.render("Searching...", True, (200, 200, 200))
                    self.screen.blit(text, (20, 140))

            # TTS box
            if self.typing_mode:
                s = pygame.Surface((self.display_width, 40))
                s.set_alpha(180)
                s.fill((0,0,0))
                self.screen.blit(s, (0, self.display_height-40))
                text = self.font.render(f"Say: {self.chat_message}", True, (255,255,255))
                self.screen.blit(text, (10, self.display_height-30))

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

        print("🛑 Shutting down...")
        self.motion.stopMove()
        self.music.stop()
        if self.audio_device:
            try: self.audio_device.unsubscribe(self.audio_service_name)
            except: pass
        if getattr(self, "sound_receiver", None):
            self.sound_receiver.close()
        if hasattr(self, 'video') and self.video_client:
            try: self.video.unsubscribe(self.video_client)
            except: pass
        self.motion.rest()
        pygame.quit()


if __name__ == "__main__":
    MIC_CHANNELS = {"left": 1, "right": 2, "front": 3, "rear": 4}

    parser = argparse.ArgumentParser()
    parser.add_argument("--ip", type=str, default="127.0.0.1")
    parser.add_argument("--port", type=int, default=9559)
    parser.add_argument("--audio-device-index", type=int, default=None)
    parser.add_argument("--mic-channel", choices=MIC_CHANNELS.keys(), default="left")
    parser.add_argument("--mic-gain", type=float, default=8.0)
    parser.add_argument("--video-scale", type=float, default=2.0)
    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,
              audio_device_index=args.audio_device_index,
              mic_channel=MIC_CHANNELS[args.mic_channel],
              mic_gain=args.mic_gain,
              video_scale=args.video_scale).run()