#!/usr/bin/env python3
"""
TEMPLATE — G1 Audio Script

Copy-paste starter demonstrating every reliable primitive for G1 audio:

    1. Auto-detect and prep the Hollyland mic   (g1_audio_devices helpers)
    2. Record N seconds via parec               (the method that actually works)
    3. Load / generate PCM (16 kHz mono int16)
    4. Play on the G1 speaker in ONE call       (_CallRequestWithParamAndBin)
    5. Always STOP the stream at the end        (mandatory reset)

Do NOT use:
    - PlayStream (chunked — unreliable, see voice_note.txt)
    - PyAudio input against the pulse device on the robot (gives silence)
    - aplay / paplay to the built-in speaker (no physical path)

Run on the robot, gemini conda env (needs _CallRequestWithParamAndBin):
    /home/unitree/miniconda3/envs/gemini/bin/python3 template_audio_script.py

Run on the workstation (edit MODE below to "tone" — no robot needed).
"""

from __future__ import annotations

import argparse
import json
import math
import struct
import subprocess
import sys
import time
from pathlib import Path

# ─── Constants ──────────────────────────────────────────────────────────────
TARGET_RATE = 16000
BIT_DEPTH = 16
CHANNELS = 1
APP_NAME = "my_app"     # pick any unique label for YOUR script
DDS_IFACE = "eth0"      # on workstation use "enp3s0"
VOLUME = 100


# ─── 1. Mic prep (optional — skip if recording is handled elsewhere) ────────
def prep_mic(keywords=("hollyland", "wireless_microphone")):
    """Auto-detect and unmute the first PulseAudio source matching keywords."""
    rc = subprocess.run(["pactl", "list", "short", "sources"], capture_output=True, text=True)
    for line in rc.stdout.splitlines():
        parts = line.split("\t")
        if len(parts) < 2:
            continue
        idx, name = parts[0], parts[1]
        if any(k in name.lower() for k in keywords):
            subprocess.run(["pactl", "set-default-source", name])
            subprocess.run(["pactl", "set-source-mute", idx, "0"])
            subprocess.run(["pactl", "set-source-volume", idx, "100%"])
            return int(idx), name
    return None, None


# ─── 2. Record via parec (the method that works on the robot) ───────────────
def record(seconds: float, source_index: int) -> bytes:
    proc = subprocess.Popen(
        ["parec", "-d", str(source_index),
         "--format=s16le", "--rate=16000", "--channels=1", "--raw"],
        stdout=subprocess.PIPE,
    )
    time.sleep(seconds)
    proc.terminate()
    return proc.stdout.read() if proc.stdout else b""


# ─── 3. Generate a demo PCM (used when MODE=tone) ───────────────────────────
def sine_tone(freq_hz: float, seconds: float) -> bytes:
    n = int(TARGET_RATE * seconds)
    out = bytearray()
    for i in range(n):
        val = int(0.3 * 32767 * math.sin(2 * math.pi * freq_hz * i / TARGET_RATE))
        out += struct.pack("<h", val)
    return bytes(out)


# ─── 4. Play PCM on the G1 speaker in ONE call ──────────────────────────────
def play_on_g1(pcm: bytes):
    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,
    )

    ChannelFactoryInitialize(0, DDS_IFACE)
    c = AudioClient()
    c.SetTimeout(10.0)
    c.Init()
    c.SetVolume(VOLUME)

    # 4a. Reset any previous stream first.
    c._Call(ROBOT_API_ID_AUDIO_STOP_PLAY, json.dumps({"app_name": APP_NAME}))
    time.sleep(0.3)

    # 4b. Unique stream_id per play.
    sid = f"s_{int(time.time() * 1000)}"
    param = json.dumps({
        "app_name": APP_NAME,
        "stream_id": sid,
        "sample_rate": TARGET_RATE,
        "channels": CHANNELS,
        "bits_per_sample": BIT_DEPTH,
    })

    # 4c. Send ALL audio in ONE call.
    duration = len(pcm) / (TARGET_RATE * BIT_DEPTH // 8)
    print(f"Playing {duration:.1f}s ...")
    c._CallRequestWithParamAndBin(ROBOT_API_ID_AUDIO_START_PLAY, param, list(pcm))
    time.sleep(duration + 0.5)

    # 4d. Stop = mandatory cleanup.
    c._Call(ROBOT_API_ID_AUDIO_STOP_PLAY, json.dumps({"app_name": APP_NAME}))


# ─── Main: pick a mode ──────────────────────────────────────────────────────
def main() -> int:
    p = argparse.ArgumentParser(description="G1 audio template")
    p.add_argument("mode", choices=["tone", "echo", "prep-mic"],
                   help="tone=440Hz beep; echo=record+play; prep-mic=auto-detect mic only")
    p.add_argument("--seconds", "-s", type=float, default=3.0,
                   help="Recording / tone duration (default 3)")
    args = p.parse_args()

    if args.mode == "prep-mic":
        idx, name = prep_mic()
        if idx is None:
            print("No wireless mic found.")
            return 1
        print(f"Prepped: [{idx}] {name}")
        return 0

    if args.mode == "tone":
        pcm = sine_tone(440.0, args.seconds)
        play_on_g1(pcm)
        return 0

    if args.mode == "echo":
        idx, name = prep_mic()
        if idx is None:
            print("No wireless mic found.")
            return 1
        print(f"Mic: [{idx}] {name}")
        print(f"Recording {args.seconds:.1f}s ... speak now")
        pcm = record(args.seconds, idx)
        if not pcm:
            print("No audio captured.")
            return 2
        print(f"Captured {len(pcm):,} bytes. Playing back on G1 ...")
        play_on_g1(pcm)
        return 0

    return 1


if __name__ == "__main__":
    sys.exit(main())