robot.py

from __future__ import annotations

import html
import math
import os
import threading
import time
import uuid
import wave
from dataclasses import dataclass
from pathlib import Path

import numpy as np


# Where captured utterances are written before the engine reads them back.
CONV_AUDIO_DIR = Path(os.getenv("REACHY_BRIDGE_CACHE", "runtime_cache")) / "conv"


AVATAR_STATES = {
    "idle",
    "listening",
    "thinking",
    "speaking",
    "nodding",
    "text_only",
    "success",
    "encourage",
}


@dataclass(slots=True)
class AvatarSnapshot:
    state: str
    title: str
    detail: str
    badge: str
    connected: bool
    connection_label: str


class ReachyBridgeRobot:
    """Drives a physical Reachy Mini when connected; otherwise a no-op (virtual).

    The hardware calls target the official `reachy_mini` SDK (see
    https://huggingface.co/docs/reachy_mini/SDK/python-sdk):
      * connect:  `with ReachyMini(connection_mode=...) as mini`
      * move:     `mini.goto_target(head=create_head_pose(...), antennas=..., body_yaw=..., method="minjerk")`
      * speak:    `mini.media.start_playing()` then `mini.media.push_audio_sample(samples)`

    Every hardware call is feature-detected so an older/newer SDK degrades to the
    virtual experience with a readable status message instead of crashing.
    """

    def __init__(self, connected: bool | None = None) -> None:
        if connected is None:
            connected = os.getenv("ROBOT_CONNECTED", "0") == "1"
        self.connected_requested = connected
        self.connection_mode = os.getenv("REACHY_CONNECTION_MODE")  # localhost_only | network
        self.media_backend = os.getenv("REACHY_MEDIA_BACKEND", "default")
        self._mini = None
        self._create_head_pose = None
        self._sdk_error: str | None = None
        self._hardware_ready = False
        self._playback_ready = False
        # Live Conversation (robot-only) state
        self._motors_enabled = False
        self._hw_lock = threading.RLock()  # serializes ALL SDK calls across threads
        self._speaking = threading.Event()  # set while the robot talks (echo guard + idle pause)
        self._idle_thread: threading.Thread | None = None
        self._idle_stop: threading.Event | None = None
        self.vad_energy = float(os.getenv("CONV_VAD_ENERGY", "0.012"))
        # NOTE: reachy_mini antenna lists are [right, left] per the SDK.

    # ------------------------------------------------------------------
    # Avatar snapshots (drive the on-screen 3D Reachy state machine)
    # ------------------------------------------------------------------
    def snapshot(self, state: str, title: str, detail: str, badge: str) -> AvatarSnapshot:
        normalized = state if state in AVATAR_STATES else "idle"
        return AvatarSnapshot(
            state=normalized,
            title=title,
            detail=detail,
            badge=badge,
            connected=self._hardware_ready,
            connection_label=self._connection_label(),
        )

    def idle_snapshot(self) -> AvatarSnapshot:
        return self.snapshot("idle", "Reachy is ready", "Pick a language route, then talk or type to start.", "Standby")

    def listening_snapshot(self, pair_label: str) -> AvatarSnapshot:
        return self.snapshot("listening", "Listening", f"Capturing the speaker for {pair_label}.", "Input Live")

    def thinking_snapshot(self) -> AvatarSnapshot:
        return self.snapshot("thinking", "Interpreting", "Understanding speech and composing the translation.", "Processing")

    def speaking_snapshot(self, text: str) -> AvatarSnapshot:
        return self.snapshot("speaking", "Reachy speaks the translation", text, "Voice Output")

    def nodding_snapshot(self, text: str) -> AvatarSnapshot:
        return self.snapshot("nodding", "Reachy confirms", text, "Motion Cue")

    def text_only_snapshot(self, text: str) -> AvatarSnapshot:
        return self.snapshot("text_only", "Translation ready", text, "Screen Output")

    def success_snapshot(self, text: str) -> AvatarSnapshot:
        return self.snapshot("success", "Great job", text, "Tutor Win")

    def encourage_snapshot(self, text: str) -> AvatarSnapshot:
        return self.snapshot("encourage", "Try once more", text, "Coach Mode")

    # ------------------------------------------------------------------
    # Turn behaviors
    # ------------------------------------------------------------------
    def perform_translation(
        self,
        text: str,
        audio_path: str | None,
        speech_supported: bool,
        turn_hint: str = "center",
    ) -> AvatarSnapshot:
        snapshot = self.speaking_snapshot(text) if speech_supported else self.text_only_snapshot(text)
        if not self._ensure_connected():
            return snapshot
        try:
            self._turn_to_speaker(turn_hint)
            if speech_supported and audio_path:
                self._play_audio(Path(audio_path))
            self._nod_once(turn_hint)
        except Exception as exc:  # never let a hardware hiccup break the turn
            self._sdk_error = str(exc)
        return snapshot

    def perform_tutor_feedback(self, success: bool, audio_path: str | None, text: str) -> AvatarSnapshot:
        snapshot = self.success_snapshot(text) if success else self.encourage_snapshot(text)
        if not self._ensure_connected():
            return snapshot
        try:
            if audio_path:
                self._play_audio(Path(audio_path))
            self._nod_once("center")
        except Exception as exc:
            self._sdk_error = str(exc)
        return snapshot

    def self_test(self) -> tuple[bool, str]:
        """One-click hardware check: connect, wiggle antennas, recenter. Safe + small."""
        if not self.connected_requested:
            return False, "Robot mode is off (virtual). Set ROBOT_CONNECTED=1 and run locally with your Reachy Mini."
        if not self._ensure_connected():
            return False, self._sdk_error or "Could not reach a Reachy Mini. Is the daemon running?"
        try:
            self._move(antennas_deg=[35, 35], duration=0.5)
            self._move(antennas_deg=[-10, -10], duration=0.5)
            self._move(head_kwargs={"yaw": 0, "pitch": 0}, antennas_deg=[0, 0], duration=0.5)
            return True, "Reachy Mini responded: antennas + head moved. Hardware looks good."
        except Exception as exc:
            self._sdk_error = str(exc)
            return False, f"Connected, but a motion call failed: {exc}"

    def connection_report(self) -> str:
        if not self.connected_requested:
            return "Virtual only (browser)"
        if self._hardware_ready:
            return "Reachy Mini linked"
        if self._sdk_error:
            return f"Robot requested, fallback active ({self._sdk_error})"
        return "Robot requested (not yet linked)"

    # ------------------------------------------------------------------
    # Hardware plumbing (feature-detected against the reachy_mini SDK)
    # ------------------------------------------------------------------
    def _connection_label(self) -> str:
        if self.connected_requested and self._hardware_ready:
            return "Robot Linked"
        if self.connected_requested:
            return "Robot Fallback"
        return "Virtual Reachy"

    def _ensure_connected(self) -> bool:
        if not self.connected_requested:
            return False
        if self._hardware_ready and self._mini is not None:
            return True

        try:
            from reachy_mini import ReachyMini
            from reachy_mini.utils import create_head_pose
        except ImportError as exc:
            self._sdk_error = f"reachy_mini SDK not installed: {exc}"
            return False

        try:
            self._create_head_pose = create_head_pose
            kwargs = {"media_backend": self.media_backend}
            if self.connection_mode:
                kwargs["connection_mode"] = self.connection_mode
            # ReachyMini() auto-detects USB/localhost vs network; connection_mode is optional.
            self._mini = self._construct_mini(ReachyMini, kwargs)
            if hasattr(self._mini, "__enter__"):
                self._mini.__enter__()
            self._enable_motors()
            self._init_playback()
            self._hardware_ready = True
            return True
        except Exception as exc:
            self._mini = None
            self._hardware_ready = False
            self._sdk_error = str(exc)
            return False

    @staticmethod
    def _construct_mini(reachy_cls, kwargs: dict):
        # Be tolerant of SDK versions that don't accept every kwarg.
        try:
            return reachy_cls(**kwargs)
        except TypeError:
            slim = {k: v for k, v in kwargs.items() if k == "connection_mode"}
            try:
                return reachy_cls(**slim)
            except TypeError:
                return reachy_cls()

    def _init_playback(self) -> None:
        media = getattr(self._mini, "media", None)
        if media is not None and hasattr(media, "start_playing"):
            try:
                media.start_playing()
                self._playback_ready = True
            except Exception as exc:
                self._sdk_error = f"Robot speaker unavailable: {exc}"

    def _enable_motors(self) -> None:
        # set_target requires motors enabled first; enable_motors() pins targets
        # to the present pose, so it must run once right after connecting.
        mini = self._mini
        if mini is None or self._motors_enabled or not hasattr(mini, "enable_motors"):
            return
        try:
            mini.enable_motors()
            self._motors_enabled = True
        except Exception as exc:
            self._sdk_error = f"enable_motors failed: {exc}"

    def _move(
        self,
        head_kwargs: dict | None = None,
        antennas_deg: list[float] | None = None,
        body_yaw_deg: float | None = None,
        duration: float = 0.6,
    ) -> None:
        if self._mini is None or not hasattr(self._mini, "goto_target"):
            return
        target: dict = {"duration": duration}
        if head_kwargs is not None and self._create_head_pose is not None:
            target["head"] = self._create_head_pose(degrees=True, mm=True, **head_kwargs)
        if antennas_deg is not None:
            target["antennas"] = np.deg2rad(np.asarray(antennas_deg, dtype=float))
        if body_yaw_deg is not None:
            target["body_yaw"] = float(np.deg2rad(body_yaw_deg))
        with self._hw_lock:
            try:
                self._mini.goto_target(method="minjerk", **target)
            except TypeError:
                target.pop("method", None)
                self._mini.goto_target(**target)

    def _turn_to_speaker(self, turn_hint: str) -> None:
        yaw_lookup = {"left": 22, "right": -22, "center": 0}
        body_yaw = yaw_lookup.get(turn_hint, 0)
        self._move(
            head_kwargs={"yaw": body_yaw * 0.4, "pitch": -4},
            antennas_deg=[12, -12],
            body_yaw_deg=body_yaw,
            duration=0.8,
        )

    def _nod_once(self, turn_hint: str) -> None:
        yaw_lookup = {"left": 9, "right": -9, "center": 0}
        yaw = yaw_lookup.get(turn_hint, 0)
        self._move(head_kwargs={"yaw": yaw, "pitch": 10}, antennas_deg=[8, -8], duration=0.35)
        time.sleep(0.2)
        self._move(head_kwargs={"yaw": yaw, "pitch": -6}, antennas_deg=[4, -4], duration=0.35)
        time.sleep(0.2)

    def _play_audio(self, audio_path: Path, expressive: bool = False, stop_event=None) -> None:
        media = getattr(self._mini, "media", None)
        if media is None or not hasattr(media, "push_audio_sample"):
            raise RuntimeError("This Reachy Mini SDK build has no media.push_audio_sample for speaker output.")
        if not self._playback_ready:
            self._init_playback()

        samples, source_rate = _read_wav_mono_float32(audio_path)
        target_rate = source_rate
        if hasattr(media, "get_output_audio_samplerate"):
            try:
                target_rate = int(media.get_output_audio_samplerate())
            except Exception:
                target_rate = source_rate
        if target_rate != source_rate:
            samples = _resample(samples, source_rate, target_rate)

        # push expects shape (samples, channels), float32
        frame = samples.reshape(-1, 1).astype(np.float32)
        with self._hw_lock:
            media.push_audio_sample(frame)
        # push_audio_sample is non-blocking; wait out the clip so motion lines up.
        clip = min(15.0, len(samples) / max(1, target_rate))
        if expressive:
            self._talk_motion(clip, stop_event)
        else:
            self._sleep_interruptible(clip, stop_event)

    # ------------------------------------------------------------------
    # Live Conversation (robot-only, hands-free) — all additive. None of this
    # runs unless start_conversation() succeeds with a linked Reachy Mini.
    # ------------------------------------------------------------------
    def conversation_ready(self) -> bool:
        if not self._ensure_connected():
            return False
        media = getattr(self._mini, "media", None)
        return media is not None and hasattr(media, "get_audio_sample")

    def start_conversation(self) -> bool:
        if not self._ensure_connected():
            return False
        media = getattr(self._mini, "media", None)
        if media is None or not hasattr(media, "get_audio_sample"):
            self._sdk_error = "This Reachy Mini SDK build has no microphone capture (media.get_audio_sample)."
            return False
        try:
            if hasattr(media, "start_recording"):
                media.start_recording()
        except Exception as exc:
            self._sdk_error = f"Could not start the robot microphone: {exc}"
            return False
        self._init_playback()
        self.start_idle_breathing()
        return True

    def stop_conversation(self) -> None:
        self.stop_idle_breathing()
        media = getattr(self._mini, "media", None)
        try:
            if media is not None and hasattr(media, "stop_recording"):
                media.stop_recording()
        except Exception:
            pass
        try:  # cut any audio still playing so Stop is immediate; re-init on next start
            if media is not None and hasattr(media, "stop_playing"):
                media.stop_playing()
                self._playback_ready = False
        except Exception:
            pass
        try:  # recenter politely
            self._move(head_kwargs={"yaw": 0, "pitch": 0}, antennas_deg=[0, 0], body_yaw_deg=0, duration=0.5)
        except Exception:
            pass

    def start_idle_breathing(self) -> None:
        if self._mini is None or not hasattr(self._mini, "set_target"):
            return
        if self._idle_thread is not None and self._idle_thread.is_alive():
            return
        self._idle_stop = threading.Event()
        self._idle_thread = threading.Thread(target=self._idle_loop, daemon=True)
        self._idle_thread.start()

    def stop_idle_breathing(self) -> None:
        if self._idle_stop is not None:
            self._idle_stop.set()
        thread = self._idle_thread
        if thread is not None and thread.is_alive():
            thread.join(timeout=1.0)
        self._idle_thread = None
        self._idle_stop = None

    def _idle_loop(self) -> None:
        """Subtle 'breathing' via instant set_target; pauses while the robot talks."""
        t0 = time.time()
        stop = self._idle_stop
        while stop is not None and not stop.is_set():
            if self._speaking.is_set():
                time.sleep(0.06)
                continue
            t = time.time() - t0
            yaw = 2.0 * math.sin(2 * math.pi * 0.13 * t)
            pitch = 3.0 * math.sin(2 * math.pi * 0.25 * t)
            sway = 5.0 * math.sin(2 * math.pi * 0.5 * t)
            try:
                self._set_target(head_kwargs={"yaw": yaw, "pitch": pitch},
                                 antennas_deg=[sway, -sway], body_yaw_deg=yaw * 0.3)
            except Exception:
                break
            time.sleep(0.05)

    def capture_utterance(self, stop_event, *, max_seconds: float = 12.0,
                          silence_ms: int = 800, listen_timeout: float = 30.0):
        """Listen on the robot mic until a full utterance ends. Returns (wav_path|None, doa)."""
        media = getattr(self._mini, "media", None)
        if media is None or not hasattr(media, "get_audio_sample"):
            return None, math.pi / 2

        sample_rate = 16000
        if hasattr(media, "get_input_audio_samplerate"):
            try:
                sample_rate = int(media.get_input_audio_samplerate()) or 16000
            except Exception:
                sample_rate = 16000
        if sample_rate <= 0:
            sample_rate = 16000

        frames: list[np.ndarray] = []
        started = False
        silence_s = 0.0
        spoken_s = 0.0
        waited_s = 0.0
        last_doa = math.pi / 2
        silence_limit = silence_ms / 1000.0

        while not stop_event.is_set():
            if self._speaking.is_set():  # never capture our own voice (echo guard)
                time.sleep(0.02)
                continue
            with self._hw_lock:
                chunk = media.get_audio_sample()
            if chunk is None:
                time.sleep(0.02)
                if not started:
                    waited_s += 0.02
                    if waited_s >= listen_timeout:
                        return None, last_doa
                continue
            chunk = np.asarray(chunk, dtype=np.float32)
            mono = chunk.mean(axis=1) if chunk.ndim > 1 else chunk
            seg = len(mono) / sample_rate

            speech = None
            if hasattr(media, "get_DoA"):
                with self._hw_lock:
                    doa = media.get_DoA()
                if doa is not None:  # MUST guard: get_DoA() can return None
                    try:
                        if isinstance(doa, (tuple, list)) and len(doa) >= 2:
                            last_doa = float(doa[0])
                            speech = bool(doa[1])
                        elif isinstance(doa, (int, float)):
                            last_doa = float(doa)  # angle only -> use energy VAD
                        else:
                            self._sdk_error = f"Unexpected get_DoA() shape: {type(doa)}"
                    except Exception as exc:
                        self._sdk_error = f"get_DoA parse failed: {exc}"
            if speech is None:  # energy fallback when DoA unavailable
                rms = float(np.sqrt(np.mean(np.square(mono)))) if mono.size else 0.0
                speech = rms > self.vad_energy

            if speech:
                started = True
                frames.append(mono)
                silence_s = 0.0
                spoken_s += seg
            elif started:
                frames.append(mono)
                silence_s += seg
                spoken_s += seg
                if silence_s >= silence_limit:
                    break
            if started and spoken_s >= max_seconds:
                break

        if stop_event.is_set() or not frames:
            return None, last_doa
        audio = np.concatenate(frames).astype(np.float32)
        if audio.size < int(0.3 * sample_rate):  # ignore sub-300ms blips
            return None, last_doa
        return _write_conversation_wav(audio, sample_rate), last_doa

    def look_at_doa(self, doa: float) -> None:
        # SDK DoA mapping: 0 = left, pi/2 = front/back, pi = right. We treat pi/2 as
        # zero yaw (facing forward) by design and only track the left/right half.
        # Instant set_target (not goto_target) so the turn never blocks Stop.
        frac = (math.pi / 2 - float(doa)) / (math.pi / 2)
        frac = max(-1.0, min(1.0, frac))
        body_yaw = frac * 18.0
        # antennas are [right, left]; perk the antenna on the speaker's side.
        antennas = [16, 4] if frac >= 0 else [4, 16]
        self._set_target(head_kwargs={"yaw": body_yaw * 0.4, "pitch": -3},
                         antennas_deg=antennas, body_yaw_deg=body_yaw)

    def perform_conversation_turn(self, audio_path, doa: float, stop_event=None) -> None:
        if stop_event is not None and stop_event.is_set():
            return
        self._speaking.set()
        try:
            self.look_at_doa(doa)
            if audio_path and not (stop_event is not None and stop_event.is_set()):
                self._play_audio(Path(audio_path), expressive=True, stop_event=stop_event)
            if not (stop_event is not None and stop_event.is_set()):
                self._excited_landed(stop_event)
        except Exception as exc:
            self._sdk_error = str(exc)
        finally:
            self._speaking.clear()

    def speak_chatter(self, audio_path, stop_event=None) -> None:
        """Play a NON-translation interjection (greeting/sign-off) on the robot speaker."""
        if not audio_path:
            return
        if stop_event is not None and stop_event.is_set():
            return
        self._speaking.set()
        try:
            self._play_audio(Path(audio_path), expressive=True, stop_event=stop_event)
        except Exception as exc:
            self._sdk_error = str(exc)
        finally:
            self._speaking.clear()

    def _excited_landed(self, stop_event=None) -> None:
        # Instant set_target wiggle so it never blocks Stop (was blocking goto_target).
        self._set_target(head_kwargs={"yaw": 0, "pitch": 6}, antennas_deg=[30, 30], body_yaw_deg=8)
        self._sleep_interruptible(0.2, stop_event)
        self._set_target(head_kwargs={"yaw": 0, "pitch": 0}, antennas_deg=[0, 0], body_yaw_deg=0)

    def _talk_motion(self, duration: float, stop_event=None) -> None:
        """Small head bobs synced to playback length (instant set_target loop)."""
        t0 = time.time()
        while time.time() - t0 < duration:
            if stop_event is not None and stop_event.is_set():
                return
            t = time.time() - t0
            bob = 6.0 * max(0.0, math.sin(t * 9.0))
            try:
                self._set_target(head_kwargs={"yaw": 3.0 * math.sin(t * 2.0), "pitch": 2.0 - bob * 0.4},
                                 antennas_deg=[6, 6])
            except Exception:
                self._sleep_interruptible(max(0.0, duration - (time.time() - t0)), stop_event)
                return
            time.sleep(0.04)

    @staticmethod
    def _sleep_interruptible(duration: float, stop_event=None) -> None:
        if stop_event is None:
            time.sleep(max(0.0, duration))
            return
        end = time.time() + max(0.0, duration)
        while time.time() < end and not stop_event.is_set():
            time.sleep(0.05)

    def _set_target(self, head_kwargs: dict | None = None, antennas_deg: list[float] | None = None,
                    body_yaw_deg: float | None = None) -> None:
        mini = self._mini
        if mini is None or not hasattr(mini, "set_target"):
            return
        kwargs: dict = {}
        # create_head_pose defaults to meters; pass mm=True to match _move's convention.
        if head_kwargs is not None and self._create_head_pose is not None:
            kwargs["head"] = self._create_head_pose(degrees=True, mm=True, **head_kwargs)
        if antennas_deg is not None:
            kwargs["antennas"] = np.deg2rad(np.asarray(antennas_deg, dtype=float))
        if body_yaw_deg is not None:
            kwargs["body_yaw"] = float(np.deg2rad(body_yaw_deg))
        if not kwargs:
            return
        try:
            with self._hw_lock:
                mini.set_target(**kwargs)
        except Exception:
            pass


def _read_wav_mono_float32(audio_path: Path) -> tuple[np.ndarray, int]:
    with wave.open(str(audio_path), "rb") as wav_file:
        channels = wav_file.getnchannels()
        sample_width = wav_file.getsampwidth()
        frame_rate = wav_file.getframerate()
        raw = wav_file.readframes(wav_file.getnframes())
    if sample_width == 2:
        audio = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
    elif sample_width == 4:
        audio = np.frombuffer(raw, dtype="<i4").astype(np.float32) / 2147483648.0
    else:
        raise RuntimeError(f"Unsupported WAV sample width: {sample_width}")
    if channels > 1:
        audio = audio.reshape(-1, channels).mean(axis=1)
    return np.clip(audio, -1.0, 1.0).astype(np.float32), frame_rate


def _resample(audio: np.ndarray, source_rate: int, target_rate: int) -> np.ndarray:
    if source_rate == target_rate or audio.size == 0:
        return audio.astype(np.float32)
    duration = audio.shape[0] / source_rate
    target_size = max(1, int(round(duration * target_rate)))
    source_positions = np.linspace(0.0, duration, num=audio.shape[0], endpoint=False)
    target_positions = np.linspace(0.0, duration, num=target_size, endpoint=False)
    return np.interp(target_positions, source_positions, audio).astype(np.float32)


def _write_conversation_wav(audio: np.ndarray, sample_rate: int) -> str:
    CONV_AUDIO_DIR.mkdir(parents=True, exist_ok=True)
    output_path = CONV_AUDIO_DIR / f"utt_{uuid.uuid4().hex}.wav"
    pcm = (np.clip(audio, -1.0, 1.0) * 32767.0).astype("<i2")
    with wave.open(str(output_path), "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(2)
        wav_file.setframerate(sample_rate)
        wav_file.writeframes(pcm.tobytes())
    return str(output_path)


def render_avatar_html(snapshot: AvatarSnapshot) -> str:
    """Emit a small state signal the 3D avatar stage reacts to (no scene reload)."""
    state = html.escape(snapshot.state)
    title = html.escape(snapshot.title)
    detail = html.escape(snapshot.detail)
    badge = html.escape(snapshot.badge)
    connection = html.escape(snapshot.connection_label)
    return f"""
    <div id="reachy-state-signal" class="reachy-state-signal" data-avatar-state="{state}" data-connection="{connection}">
      <div class="reachy-state-head">
        <span class="reachy-state-badge">{badge}</span>
        <span class="reachy-state-conn">{connection}</span>
      </div>
      <div class="reachy-state-copy">
        <strong>{title}</strong>
        <span>{detail}</span>
      </div>
    </div>
    """