#!/usr/bin/env python3
"""
Test script for Marcus env — TTS + Audio playback on G1 built-in speaker.
Deploy to robot: scp test_marcus_audio.py unitree@192.168.123.164:/home/unitree/Marcus/
Run on robot:    conda activate marcus && python3 ~/Marcus/test_marcus_audio.py
"""

import sys
import time
import json
import struct
import math

def check_sdk():
    """Check if marcus env has the required SDK methods."""
    print("=== SDK CHECK ===")
    try:
        from unitree_sdk2py.g1.audio.g1_audio_client import AudioClient
        from unitree_sdk2py.rpc.client import Client
        print("  unitree_sdk2py: installed")

        has_bin = hasattr(Client, '_CallRequestWithParamAndBin')
        print(f"  _CallRequestWithParamAndBin: {'YES' if has_bin else 'NO (need upgrade: pip install --upgrade unitree_sdk2py)'}")
        return has_bin
    except ImportError:
        print("  unitree_sdk2py: NOT INSTALLED")
        print("  Run: pip install unitree_sdk2py")
        return False


def test_tts():
    """Test TTS on G1 speaker."""
    from unitree_sdk2py.core.channel import ChannelFactoryInitialize
    from unitree_sdk2py.g1.audio.g1_audio_client import AudioClient

    print("\n=== TTS TEST ===")
    ChannelFactoryInitialize(0, "eth0")
    c = AudioClient()
    c.SetTimeout(10.0)
    c.Init()

    # Volume
    code, vol = c.GetVolume()
    print(f"  Volume: {vol}")
    c.SetVolume(100)

    # LED
    print("  LED -> Green")
    c.LedControl(0, 255, 0)
    time.sleep(1)
    c.LedControl(0, 0, 0)

    # TTS English
    print("  TTS: 'Hello, Marcus is ready'")
    c.TtsMaker("Hello, Marcus is ready", 0)
    time.sleep(4)

    print("  TTS test done.")


def test_audio_playback():
    """Test audio playback with a sine tone."""
    from unitree_sdk2py.core.channel import ChannelFactoryInitialize
    from unitree_sdk2py.g1.audio.g1_audio_client import AudioClient
    from unitree_sdk2py.g1.audio.g1_audio_api import (
        ROBOT_API_ID_AUDIO_START_PLAY,
        ROBOT_API_ID_AUDIO_STOP_PLAY,
    )

    print("\n=== AUDIO PLAYBACK TEST ===")
    ChannelFactoryInitialize(0, "eth0")
    c = AudioClient()
    c.SetTimeout(10.0)
    c.Init()
    c.SetVolume(100)

    # Generate 2s sine tone at 440Hz
    rate = 16000
    duration = 2
    pcm = b""
    for i in range(rate * duration):
        val = int(30000 * math.sin(2 * math.pi * 440 * i / rate))
        pcm += struct.pack("<h", val)

    print(f"  Sine tone: {len(pcm)} bytes, {rate}Hz, {duration}s")

    # Stop previous, send all in one call with unique stream_id
    c._Call(ROBOT_API_ID_AUDIO_STOP_PLAY, json.dumps({"app_name": "test"}))
    time.sleep(0.3)

    sid = f"s_{int(time.time() * 1000)}"
    param = json.dumps({
        "app_name": "test",
        "stream_id": sid,
        "sample_rate": 16000,
        "channels": 1,
        "bits_per_sample": 16,
    })

    print("  Playing tone...")
    result = c._CallRequestWithParamAndBin(ROBOT_API_ID_AUDIO_START_PLAY, param, list(pcm))
    print(f"  Result: {result}")
    time.sleep(duration + 1)

    c._Call(ROBOT_API_ID_AUDIO_STOP_PLAY, json.dumps({"app_name": "test"}))
    print("  Playback test done.")


def main():
    has_sdk = check_sdk()

    if not has_sdk:
        print("\nSDK missing or outdated. Fix with:")
        print("  conda activate marcus")
        print("  pip install --upgrade unitree_sdk2py")
        sys.exit(1)

    test_tts()
    test_audio_playback()

    print("\n=== ALL TESTS DONE ===")
    print("Did you hear:")
    print("  1. LED flash green?")
    print("  2. 'Hello, Marcus is ready' (TTS)?")
    print("  3. 2-second beep tone (audio playback)?")


if __name__ == "__main__":
    main()