robot_controller.py

"""Headless robot controller for THAU Reachy."""

from __future__ import annotations

import logging
import math
import os
import shutil
import tempfile
import threading
import time
from pathlib import Path
from typing import Any

import soundfile as sf


LOGGER = logging.getLogger("thau_reachy.robot")


class HeadlessRobotController:
    def __init__(self) -> None:
        self.state = "disconnected"
        self.detail = "Pendiente"
        self.robot: Any = None
        self.media: Any = None
        self.no_robot = os.getenv("THAU_NO_ROBOT", "0").lower() in {"1", "true", "yes"}
        self.simulation = os.getenv("REACHY_SIM", "0").lower() in {"1", "true", "yes"}
        self._pc_stream: Any = None
        self._pc_source = "Micrófono predeterminado del PC"
        self._recording_sample_rate = 16_000
        self._recording = False
        self._recording_stop = threading.Event()
        self._recording_thread: threading.Thread | None = None
        self._audio_chunks: list[Any] = []
        self._audio_lock = threading.Lock()
        self._connect_lock = threading.Lock()
        self._command_lock = threading.Lock()
        self._speaking_lock = threading.Lock()
        self._speaking = False
        self._external_robot = False

    def connect(self) -> None:
        with self._connect_lock:
            if self.state == "ready":
                return
            if self.no_robot:
                self.state = "ready"
                self.detail = "Modo sin robot activo"
                LOGGER.info("Modo sin robot activo")
                return
            self.state = "connecting"
            self.detail = "Conectando con Reachy Mini"
            try:
                from reachy_mini import ReachyMini
                from reachy_mini.media.media_manager import MediaBackend, MediaManager

                self.robot = ReachyMini(
                    connection_mode="localhost_only",
                    media_backend="no_media",
                    timeout=8.0,
                    log_level="WARNING",
                )
                if not self.simulation:
                    self.media = MediaManager(backend=MediaBackend.LOCAL, log_level="WARNING")
                    self.robot.media_manager = self.media
                self.robot.enable_motors()
                if self.media is not None:
                    try:
                        self.robot.enable_wobbling()
                    except Exception as exc:
                        LOGGER.warning("No se pudo habilitar wobbling: %s", exc)
                self.state = "ready"
                self.detail = "Robot listo"
                LOGGER.info("Robot listo")
            except Exception as exc:
                self.state = "error"
                self.detail = f"{type(exc).__name__}: {exc}"
                self.robot = None
                self.media = None
                LOGGER.exception("No se pudo conectar el robot")

    def close(self) -> None:
        self._recording_stop.set()
        if self._recording_thread:
            self._recording_thread.join(timeout=1.0)
        try:
            if self.media is not None and not self._external_robot:
                self.media.stop_recording()
                self.media.stop_playing()
                self.media.close()
        except Exception:
            pass
        try:
            if self.robot is not None and not self._external_robot:
                self.robot.disable_wobbling()
        except Exception:
            pass

    def attach_robot(self, reachy_mini: Any) -> None:
        self.robot = reachy_mini
        self.media = getattr(reachy_mini, "media", None)
        self.state = "ready"
        self.detail = "Robot adjuntado desde ReachyMiniApp"
        self.no_robot = False
        self._external_robot = True
        try:
            self.robot.enable_motors()
        except Exception:
            pass
        if self.media is not None:
            try:
                self.robot.enable_wobbling()
            except Exception as exc:
                LOGGER.warning("No se pudo habilitar wobbling: %s", exc)

    def _require_robot(self) -> Any:
        if self.state != "ready" or self.robot is None:
            self.connect()
        if self.state != "ready" or self.robot is None:
            raise RuntimeError(self.detail)
        return self.robot

    def motion(self, name: str) -> None:
        if self.no_robot:
            return
        robot = self._require_robot()
        from reachy_mini.utils import create_head_pose

        poses = {
            "listening": (create_head_pose(pitch=4), [10, 10], 0.35),
            "thinking": (create_head_pose(roll=-3, pitch=2, yaw=6), [4, -4], 0.4),
            "neutral": (create_head_pose(), [0, 0], 0.7),
            "wake": (create_head_pose(pitch=3), [8, 8], 0.6),
        }
        head, antennas_deg, duration = poses[name]
        antennas = [math.radians(value) for value in antennas_deg]
        with self._command_lock:
            robot.enable_motors()
            robot.goto_target(head=head, antennas=antennas, body_yaw=0.0, duration=duration)

    def set_full_target(
        self,
        head: Any | None = None,
        antennas: Any | None = None,
        body_yaw: float | None = None,
    ) -> None:
        if self.no_robot:
            return
        robot = self._require_robot()
        with self._command_lock:
            robot.set_target(head=head, antennas=antennas, body_yaw=body_yaw)

    def look_at_image(self, u: int, v: int, duration: float = 0.2) -> None:
        if self.no_robot:
            return
        robot = self._require_robot()
        with self._command_lock:
            robot.look_at_image(u, v, duration=duration)

    def turn_toward_sound(self, angle_radians: float) -> None:
        if self.no_robot:
            return
        from face_tracking.sound_tracker import SoundTracker

        robot = self._require_robot()
        x, y = SoundTracker.angle_to_target(angle_radians)
        with self._command_lock:
            robot.look_at_world(x=x, y=y, z=0.0, duration=0.25)

    def get_direction_of_arrival(self) -> tuple[float, bool] | None:
        if self.no_robot or self.media is None:
            return None
        try:
            return self.media.get_DoA()
        except Exception:
            return None

    def get_camera_frame(self) -> Any | None:
        if self.no_robot or self.media is None:
            return None
        try:
            return self.media.get_frame()
        except Exception:
            return None

    def start_microphone(self) -> dict[str, Any]:
        if self.no_robot:
            raise RuntimeError("El modo sin robot usa el micrófono del navegador, no el del robot")
        self._require_robot()
        if self._recording:
            raise RuntimeError("Ya hay una grabación en curso")

        with self._audio_lock:
            self._audio_chunks = []
        self._recording_stop.clear()
        self._recording = True

        if self.simulation:
            try:
                import sounddevice as sd
            except Exception as exc:
                self._recording = False
                raise RuntimeError(
                    "REACHY_SIM=1 requiere sounddevice en el entorno activo"
                ) from exc

            def receive_audio(indata: Any, _frames: int, _time: Any, status: Any) -> None:
                if status:
                    LOGGER.warning("Estado del micrófono del PC: %s", status)
                if self._recording and indata.size:
                    with self._audio_lock:
                        self._audio_chunks.append(indata.copy())

            stream = sd.InputStream(
                samplerate=16_000,
                channels=1,
                dtype="float32",
                callback=receive_audio,
            )
            stream.start()
            self._pc_stream = stream
            self._recording_sample_rate = 16_000
        else:
            if self.media is None or self.media.audio is None:
                self._recording = False
                raise RuntimeError("No se pudo abrir Reachy Mini Audio")
            self._recording_sample_rate = self.media.get_input_audio_samplerate()
            self.media.start_recording()
            self._recording_thread = threading.Thread(
                target=self._recording_loop,
                name="thau-reachy-microphone",
                daemon=True,
            )
            self._recording_thread.start()

        self.motion("listening")
        return {
            "sampleRate": self._recording_sample_rate,
            "channels": 1 if self.simulation else self.media.get_input_channels(),
            "source": self._pc_source if self.simulation else "Reachy Mini Audio",
        }

    def _recording_loop(self) -> None:
        started = time.monotonic()
        while not self._recording_stop.is_set() and time.monotonic() - started < 30:
            sample = self.media.get_audio_sample()
            if sample is not None and sample.size:
                with self._audio_lock:
                    self._audio_chunks.append(sample.copy())
            else:
                time.sleep(0.005)

    def stop_microphone(self) -> Path:
        if not self._recording:
            raise RuntimeError("No había una grabación activa")

        self._recording_stop.set()
        if self._recording_thread:
            self._recording_thread.join(timeout=2.0)

        if self.simulation:
            if self._pc_stream is not None:
                self._pc_stream.stop()
                self._pc_stream.close()
                self._pc_stream = None
        elif self.media is not None:
            self.media.stop_recording()

        self._recording = False
        self.motion("thinking")

        import numpy as np

        with self._audio_lock:
            chunks = self._audio_chunks
            self._audio_chunks = []

        if not chunks:
            raise RuntimeError("No se capturaron muestras de audio del cliente")

        audio = np.concatenate(chunks, axis=0).astype(np.float32, copy=False)
        if audio.ndim == 2:
            audio = audio.mean(axis=1)
        peak = float(np.max(np.abs(audio))) if audio.size else 0.0
        if peak < 1e-4:
            raise RuntimeError("El audio está vacío o en silencio")

        output = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
        output.close()
        sf.write(output.name, audio, self._recording_sample_rate, subtype="PCM_16")
        return Path(output.name)

    def copy_client_audio(self, source: Path, destination: Path) -> None:
        destination.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(source, destination)

    def play_audio_file(self, audio_path: Path) -> None:
        if self.no_robot:
            LOGGER.info("Modo sin robot: audio TTS generado en %s", audio_path)
            return
        if self.simulation or self.media is None:
            LOGGER.info("TTS generado en %s (sin reproducción local en simulación)", audio_path)
            duration = sf.info(audio_path).duration
            with self._speaking_lock:
                self._speaking = True
            try:
                time.sleep(float(duration))
            finally:
                with self._speaking_lock:
                    self._speaking = False
            return

        duration = sf.info(audio_path).duration
        self.motion("neutral")
        with self._speaking_lock:
            self._speaking = True
        try:
            self.media.play_sound(str(audio_path))
            time.sleep(float(duration) + 0.15)
        finally:
            with self._speaking_lock:
                self._speaking = False

    def is_speaking(self) -> bool:
        with self._speaking_lock:
            return self._speaking

    def status(self) -> dict[str, Any]:
        return {
            "robotReady": self.state == "ready",
            "robotState": self.state,
            "robotDetail": self.detail,
            "simulation": self.simulation,
            "noRobot": self.no_robot,
            "cameraAvailable": self.media is not None and not self.no_robot,
            "audioAvailable": self.media is not None and not self.no_robot,
        }