storeos_robot/app.py

"""StoreOS — Decentralized AI store brain running on Reachy Mini.

Flow:
  1. Robot wakes up, greets
  2. Mic captures audio, VAD detects speech
  3. Speech → 0G Whisper (STT)
  4. Text → 0G LLM (chat completions)
  5. Response → parse actions → Kokoro TTS
  6. Audio → robot speaker + robot performs actions
  7. DoA tracking keeps head pointed at customer
"""

import asyncio
import logging
import os
import threading
import numpy as np

from reachy_mini.apps.app import ReachyMiniApp
from reachy_mini.reachy_mini import ReachyMini

from .config import ZG_API_KEY, SILENCE_THRESHOLD, SPEECH_START_FRAMES, SPEECH_END_FRAMES, SAMPLE_RATE
from . import zg_api
from . import tts
from . import notify
from . import display
from .motion import MotionController, parse_response
from .catalog import PRODUCTS
from .config import STORE_NAME, CHAIN_AUTORUN

try:
    from . import chain
    from . import inft
    from . import storage
    _chain_available = True
except ImportError as e:
    logging.getLogger(__name__).warning(f"Chain/INFT/Storage not available: {e}")
    chain = None
    inft = None
    storage = None
    _chain_available = False

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")


class StoreOSApp(ReachyMiniApp):
    """Reachy Mini store assistant powered by 0G decentralized AI."""

    custom_app_url = None
    request_media_backend = os.environ.get("REACHY_MEDIA_BACKEND", "local")

    def __init__(self):
        super().__init__()
        self._history: list[dict] = []
        self._audio_enabled = False

    def _log(self, msg):
        import sys
        line = f"STOREOS: {msg}"
        print(line, file=sys.stderr, flush=True)
        try:
            with open("/tmp/storeos.log", "a") as f:
                f.write(line + "\n")
        except: pass

    def wrapped_run(self):
        self._log("wrapped_run entered")

        display.set_store_name(STORE_NAME)
        try:
            display.start_server()
            self._log("display on :7860")
        except OSError as e:
            self._log(f"display server disabled: {e}")

        try:
            super().wrapped_run()
        except Exception as e:
            self._log(f"wrapped_run failed: {e}")
            import traceback
            traceback.print_exc(file=open("/tmp/storeos.log", "a"))
            raise

    def run(self, reachy_mini: ReachyMini, stop_event: threading.Event) -> None:
        self._log("run() called")

        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            loop.run_until_complete(self._main(reachy_mini, stop_event))
        except KeyboardInterrupt:
            pass
        except Exception as e:
            self._log(f"_main crashed: {e}")
            import traceback
            traceback.print_exc(file=open("/tmp/storeos.log", "a"))
            raise
        finally:
            loop.close()

    async def _main(self, mini: ReachyMini, stop_event: threading.Event):
        self._log("_main entered")
        self._log(f"media_backend attr = {getattr(self, 'media_backend', 'NOT SET')}")

        try:
            motion = MotionController(mini)
            await motion.start()
            self._log("motion started")
        except Exception as e:
            self._log(f"motion failed: {e}")
            return

        media_backend = (self.media_backend or "").lower()
        if media_backend == "no_media":
            self._log("audio disabled by REACHY_MEDIA_BACKEND=no_media")
        else:
            try:
                mini.media.start_recording()
                mini.media.start_playing()
                self._audio_enabled = True
                self._log("media started")
            except Exception as e:
                self._audio_enabled = False
                self._log(f"media failed: {e}")
                raise RuntimeError("Reachy media backend did not start; mic/speaker unavailable") from e

        try:
            mini.wake_up()
            self._log("robot awake")
        except Exception as e:
            self._log(f"wake_up failed: {e}")

        # 0G Chain / Storage / INFT. Wallet init is safe; startup transactions are opt-in.
        if _chain_available and chain:
            try:
                addr = chain.init_wallet()
                if addr:
                    bal = await chain.get_balance()
                    logger.info(f"0G Chain: {bal:.4f} 0G at {addr}")
                    await notify.bot_started(STORE_NAME, addr, bal)
                    if CHAIN_AUTORUN:
                        for p in PRODUCTS[:3]:
                            tx = await chain.register_product(p)
                            if tx:
                                url = chain.get_explorer_url(tx)
                                logger.info(f"Registered '{p['name']}' → {url}")
                                await notify.product_registered(p, tx, url)
                    else:
                        logger.info("Startup chain transactions disabled; set STOREOS_CHAIN_AUTORUN=true to enable")
            except Exception as e:
                logger.warning(f"Chain init failed: {e}")

        if CHAIN_AUTORUN and _chain_available and storage:
            try:
                result = await storage.upload_catalog([p for p in PRODUCTS])
                if result:
                    logger.info(f"Catalog on 0G Storage: {result['root_hash']}")
            except Exception as e:
                logger.warning(f"Storage upload failed: {e}")

        if CHAIN_AUTORUN and _chain_available and inft:
            try:
                result = await inft.mint_store_brain(STORE_NAME, self._history)
                if result and not result.get("pending"):
                    logger.info(f"INFT minted! Token #{result.get('token_id')}")
            except Exception as e:
                logger.warning(f"INFT mint failed: {e}")

        logger.info("StoreOS started — robot is ready")

        try:
            motion.trigger("greet")
        except Exception as e:
            logger.warning(f"Greet failed: {e}")

        if self._audio_enabled:
            try:
                greeting = "Welcome to the store! I can help you find anything. Just talk to me."
                await self._speak(mini, motion, greeting)
                loop = asyncio.get_event_loop()
                await self._drain_mic(mini, loop)
            except Exception as e:
                self._log(f"Greeting TTS failed: {e}")
        else:
            self._log("Audio disabled; voice greeting and listener are skipped")

        self._log("Starting main loops...")
        tasks = [
            asyncio.create_task(motion.doa_tracking_loop(), name="doa"),
            asyncio.create_task(motion.motion_loop(), name="motion"),
            asyncio.create_task(motion.action_loop(), name="actions"),
        ]
        if self._audio_enabled:
            tasks.append(asyncio.create_task(self._listen_loop(mini, motion), name="listen"))
        logger.info(f"All {len(tasks)} tasks launched — {'voice listener enabled' if self._audio_enabled else 'voice listener disabled'}")

        logger.info(f"stop_event state: {stop_event.is_set()}")
        try:
            while not stop_event.is_set():
                await asyncio.sleep(0.5)
            logger.info("stop_event was set — shutting down")
        except Exception as e:
            logger.error(f"Main loop error: {e}")
        finally:
            logger.info("Shutting down...")
            for t in tasks:
                t.cancel()
            await asyncio.gather(*tasks, return_exceptions=True)
            if self._audio_enabled:
                mini.media.stop_recording()
                mini.media.stop_playing()
            mini.goto_sleep()

    async def _listen_loop(self, mini: ReachyMini, motion: MotionController):
        """Continuous listen → transcribe → respond loop."""
        self._log("listen_loop started — waiting for speech")
        loop = asyncio.get_event_loop()

        while True:
            pcm_bytes = await self._capture_speech(mini, loop)
            if not pcm_bytes or len(pcm_bytes) < SAMPLE_RATE:  # < 0.5s
                continue
            self._log(f"got speech: {len(pcm_bytes)} bytes")

            self._log(f"captured {len(pcm_bytes) / (SAMPLE_RATE * 2):.1f}s of speech")
            motion.trigger("think")
            display.show_thinking()

            # 0G Whisper STT
            try:
                text = await zg_api.transcribe(pcm_bytes)
            except Exception as e:
                self._log(f"STT failed: {e}")
                continue

            if not text:
                self._log("STT returned empty text")
                continue

            self._log(f"User: {text}")
            await notify.customer_spoke(text)

            # 0G LLM
            try:
                raw_reply = await zg_api.chat(text, self._history)
            except Exception as e:
                self._log(f"LLM failed: {e}")
                motion.trigger("shake")
                continue

            clean, actions, sales = parse_response(raw_reply)
            self._log(f"Assistant: {clean}")
            if actions:
                self._log(f"Actions: {actions}")

            await notify.assistant_replied(clean)
            self._history.append({"role": "user", "content": text})
            self._history.append({"role": "assistant", "content": clean})

            # Show mentioned product on display tablet
            mentioned = self._find_mentioned_product(clean)
            if mentioned:
                display.show_product(mentioned, clean)
            else:
                display.show_idle(clean[:80] if clean else "How can I help?")
            if len(self._history) > 20:
                self._history = self._history[-12:]

            # Execute robot actions
            for a in actions:
                motion.trigger(a)
            if not actions:
                motion.trigger("nod")

            # Process sales on 0G Chain
            for sale_name in sales:
                product = next((p for p in PRODUCTS if p["name"].lower() == sale_name.lower()), None)
                if product:
                    product["stock"] = max(0, product["stock"] - 1)
                    if _chain_available and chain:
                        try:
                            tx = await chain.record_sale(product["name"], product["price"])
                            if tx:
                                url = chain.get_explorer_url(tx)
                                logger.info(f"SALE on-chain: {product['name']} → {url}")
                                await notify.sale_recorded(product, tx, url)
                        except Exception as e:
                            logger.warning(f"Sale chain tx failed: {e}")

            # TTS + playback
            await self._speak(mini, motion, clean)
            # Drain mic buffer so we don't hear our own TTS
            await self._drain_mic(mini, loop)

    async def _drain_mic(self, mini: ReachyMini, loop):
        """Read and discard all buffered mic audio, then wait for silence."""
        self._log("draining mic buffer...")
        for _ in range(100):
            try:
                samples = await loop.run_in_executor(None, mini.media.get_audio_sample)
                if samples is None:
                    break
            except Exception:
                break
            await asyncio.sleep(0.01)
        # Extra pause for speaker to fully stop
        await asyncio.sleep(0.5)
        self._log("mic drained")

    async def _capture_speech(self, mini: ReachyMini, loop) -> bytes:
        """VAD-based speech capture. Returns PCM16 bytes of the utterance."""
        chunks: list[bytes] = []
        speech_frames = 0
        silence_frames = 0
        recording = False
        _energy_log_counter = 0

        while True:
            try:
                samples = await loop.run_in_executor(None, mini.media.get_audio_sample)
            except Exception as e:
                if _energy_log_counter == 0:
                    self._log(f"get_audio_sample error: {e}")
                await asyncio.sleep(0.02)
                continue

            if samples is None or (hasattr(samples, 'shape') and samples.shape[0] == 0):
                await asyncio.sleep(0.01)
                continue

            if _energy_log_counter == 0:
                self._log(f"audio sample: shape={samples.shape}, dtype={samples.dtype}, range=[{samples.min():.4f}, {samples.max():.4f}]")

            if samples.ndim == 2 and samples.shape[1] >= 2:
                mono = samples.mean(axis=1)
            else:
                mono = samples.ravel()

            energy = np.abs(mono).max()

            _energy_log_counter += 1
            if _energy_log_counter % 200 == 0:
                self._log(f"mic peak={energy:.4f} threshold={SILENCE_THRESHOLD} recording={recording}")
                _energy_log_counter = 0

            if energy > SILENCE_THRESHOLD:
                speech_frames += 1
                silence_frames = 0
            else:
                silence_frames += 1
                if not recording:
                    speech_frames = 0

            if not recording and speech_frames >= SPEECH_START_FRAMES:
                recording = True
                display.show_listening()
                logger.debug("Speech started")

            if recording:
                pcm = (mono * 32767).astype(np.int16).tobytes()
                chunks.append(pcm)

                if silence_frames >= SPEECH_END_FRAMES:
                    logger.debug("Speech ended")
                    return b"".join(chunks)

                # Max 15 seconds
                total_bytes = sum(len(c) for c in chunks)
                if total_bytes > SAMPLE_RATE * 2 * 15:
                    return b"".join(chunks)

    async def _speak(self, mini: ReachyMini, motion: MotionController, text: str):
        """TTS → robot speaker."""
        if not self._audio_enabled:
            self._log("Skipping TTS — audio disabled")
            return

        self._log(f"TTS start: {text[:60]}...")
        motion.set_speaking(True)
        try:
            audio = await tts.synthesize(text)
            if audio is not None and len(audio) > 0:
                self._log(f"TTS got {len(audio)} samples ({len(audio)/SAMPLE_RATE:.1f}s), pushing to speaker")
                loop = asyncio.get_event_loop()
                audio_2d = audio.reshape(-1, 1).astype(np.float32)
                chunk_size = SAMPLE_RATE // 4
                for i in range(0, len(audio_2d), chunk_size):
                    chunk = audio_2d[i:i + chunk_size]
                    await loop.run_in_executor(None, mini.media.push_audio_sample, chunk)
                    await asyncio.sleep(0.01)
                duration = len(audio) / SAMPLE_RATE
                self._log(f"TTS playback done, waiting {duration*0.3:.1f}s")
                await asyncio.sleep(duration * 0.3)
            else:
                self._log("TTS returned no audio!")
        except Exception as e:
            self._log(f"TTS/playback failed: {e}")
        finally:
            motion.set_speaking(False)

    def _find_mentioned_product(self, text: str) -> dict | None:
        text_lower = text.lower()
        best = None
        best_len = 0
        for p in PRODUCTS:
            name_lower = p["name"].lower()
            if name_lower in text_lower and len(name_lower) > best_len:
                best = p
                best_len = len(name_lower)
        return best