#!/usr/bin/env python3
"""
G1 Interactive Voice Player
----------------------------
Lists WAV files and lets you pick which ones to play on the G1 built-in speaker.
Uses _CallRequestWithParamAndBin with format params (the working method).

Usage (run ON the robot):
  python3 g1_interactive_player.py "/home/unitree/SanadVoice/recorded voices"
  python3 g1_interactive_player.py "/home/unitree/SanadVoice/recorded voices" --volume 80
"""

import sys
import time
import wave
import json
import argparse
import numpy as np
from pathlib import Path
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,
)

TARGET_RATE = 16000
APP_NAME = "p"


def load_wav(path: str):
    with wave.open(path, "rb") as wf:
        rate = wf.getframerate()
        nch = wf.getnchannels()
        audio = np.frombuffer(wf.readframes(wf.getnframes()), dtype=np.int16)

    if nch == 2:
        audio = audio.reshape(-1, 2).mean(axis=1).astype(np.int16)

    if rate != TARGET_RATE:
        tl = int(len(audio) * TARGET_RATE / rate)
        audio = np.interp(
            np.linspace(0, len(audio), tl, endpoint=False),
            np.arange(len(audio)),
            audio.astype(np.float64),
        ).astype(np.int16)

    return audio


def play_pcm(client, pcm: np.ndarray):
    # Single call with a unique stream_id — chunked repeats with the same sid
    # are dropped by the audio service, per voice_note.txt.
    client._Call(ROBOT_API_ID_AUDIO_STOP_PLAY, json.dumps({"app_name": APP_NAME}))
    time.sleep(0.2)

    pcm_bytes = pcm.tobytes()
    duration = len(pcm) / TARGET_RATE
    sid = f"s_{int(time.time() * 1000)}"
    param = json.dumps({
        "app_name": APP_NAME,
        "stream_id": sid,
        "sample_rate": TARGET_RATE,
        "channels": 1,
        "bits_per_sample": 16,
    })

    start = time.time()
    client._CallRequestWithParamAndBin(
        ROBOT_API_ID_AUDIO_START_PLAY, param, list(pcm_bytes)
    )
    time.sleep(duration + 0.3)
    client._Call(ROBOT_API_ID_AUDIO_STOP_PLAY, json.dumps({"app_name": APP_NAME}))
    return time.time() - start


def show_menu(wavs):
    print(f"\n{'#':<4} {'Name':<45} {'Duration':>8}")
    print("-" * 59)
    for i, (p, dur) in enumerate(wavs):
        print(f"{i:<4} {p.stem:<45} {dur:>6.1f}s")
    print("-" * 59)
    print(f"{'a':<4} {'Play ALL':<45}")
    print(f"{'m':<4} {'Show menu':<45}")
    print(f"{'q':<4} {'Quit':<45}")
    print()


def main():
    parser = argparse.ArgumentParser(description="G1 Interactive Voice Player")
    parser.add_argument("directory", help="Directory containing WAV files")
    parser.add_argument("--volume", "-v", type=int, default=100, help="Volume 0-100")
    parser.add_argument("--interface", "-i", default="eth0", help="Network interface")

    args = parser.parse_args()

    wav_dir = Path(args.directory)
    if not wav_dir.is_dir():
        print(f"Not a directory: {wav_dir}")
        sys.exit(1)

    wav_files = []
    for p in sorted(wav_dir.glob("*.wav")):
        try:
            with wave.open(str(p), "rb") as wf:
                dur = wf.getnframes() / wf.getframerate()
            wav_files.append((p, dur))
        except Exception:
            pass

    if not wav_files:
        print(f"No WAV files found in {wav_dir}")
        sys.exit(1)

    print(f"Connecting to G1 speaker...")
    ChannelFactoryInitialize(0, args.interface)
    client = AudioClient()
    client.SetTimeout(10.0)
    client.Init()
    client.SetVolume(args.volume)
    print(f"Ready. Volume: {args.volume}%")

    show_menu(wav_files)

    while True:
        try:
            choice = input("Pick a number (or a/m/q): ").strip().lower()
        except (KeyboardInterrupt, EOFError):
            print("\nBye.")
            break

        if choice == "q":
            print("Bye.")
            break

        if choice == "m":
            show_menu(wav_files)
            continue

        if choice == "a":
            print(f"\nPlaying all {len(wav_files)} files...\n")
            for i, (p, dur) in enumerate(wav_files):
                print(f"  [{i+1}/{len(wav_files)}] {p.stem} ({dur:.1f}s)...", end=" ", flush=True)
                pcm = load_wav(str(p))
                elapsed = play_pcm(client, pcm)
                print(f"done ({elapsed:.1f}s)")
                time.sleep(0.5)
            print("\nAll done.\n")
            continue

        indices = []
        for part in choice.replace(" ", "").split(","):
            if "-" in part:
                try:
                    a, b = part.split("-", 1)
                    indices.extend(range(int(a), int(b) + 1))
                except ValueError:
                    pass
            else:
                try:
                    indices.append(int(part))
                except ValueError:
                    pass

        if not indices:
            print("Invalid. Enter number, range (2-5), list (1,3,5), a, m, or q.")
            continue

        for idx in indices:
            if idx < 0 or idx >= len(wav_files):
                print(f"  #{idx} out of range, skipping.")
                continue

            p, dur = wav_files[idx]
            print(f"  Playing: {p.stem} ({dur:.1f}s)...", end=" ", flush=True)
            pcm = load_wav(str(p))
            elapsed = play_pcm(client, pcm)
            print(f"done ({elapsed:.1f}s)")

        print()


if __name__ == "__main__":
    main()