#!/usr/bin/env python3
"""
G1 WAV Player — plays WAV files through the G1 built-in speaker.
Runs ON the robot via AudioClient._CallRequestWithParamAndBin.

(PlayStream chunked mode is unreliable — see voice_note.txt.)

Usage:
  python3 g1_play_wav.py <wav_file>
  python3 g1_play_wav.py <wav_file> --volume 80
  python3 g1_play_wav.py --list <directory>
  python3 g1_play_wav.py --all <directory>

Examples:
  python3 g1_play_wav.py "/home/unitree/SanadVoice/recorded voices/welcome_single.wav"
  python3 g1_play_wav.py --list "/home/unitree/SanadVoice/recorded voices"
  python3 g1_play_wav.py --all "/home/unitree/SanadVoice/recorded voices"
"""

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 = "g1_player"


def init_client(interface: str = "eth0", volume: int = 100):
    ChannelFactoryInitialize(0, interface)
    client = AudioClient()
    client.SetTimeout(10.0)
    client.Init()
    client.SetVolume(volume)
    return client


def load_wav(path: str):
    """Load WAV and convert to 16kHz mono 16-bit PCM."""
    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:
        target_len = int(len(audio) * TARGET_RATE / rate)
        audio = np.interp(
            np.linspace(0, len(audio), target_len, endpoint=False),
            np.arange(len(audio)),
            audio.astype(np.float64),
        ).astype(np.int16)

    return audio


def play_wav(client, pcm: np.ndarray, label: str = ""):
    """Send PCM to G1 speaker in ONE call (the working method)."""
    pcm_bytes = pcm.tobytes()
    duration = len(pcm) / TARGET_RATE

    if label:
        print(f"  {label} ({duration:.1f}s)...", end=" ", flush=True)
    else:
        print(f"  Playing {duration:.1f}s...", end=" ", flush=True)

    client._Call(ROBOT_API_ID_AUDIO_STOP_PLAY, json.dumps({"app_name": APP_NAME}))
    time.sleep(0.2)

    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.5)
    client._Call(ROBOT_API_ID_AUDIO_STOP_PLAY, json.dumps({"app_name": APP_NAME}))
    print(f"done ({time.time() - start:.1f}s)")


def list_wavs(directory: str):
    """List all WAV files in a directory."""
    wavs = sorted(Path(directory).glob("*.wav"))
    if not wavs:
        print(f"No WAV files in {directory}")
        return []

    print(f"\n{'#':<4} {'Name':<40} {'Duration':>10}")
    print("-" * 56)
    for i, p in enumerate(wavs):
        try:
            with wave.open(str(p), "rb") as wf:
                dur = wf.getnframes() / wf.getframerate()
            print(f"{i:<4} {p.stem:<40} {dur:>8.1f}s")
        except Exception as e:
            print(f"{i:<4} {p.stem:<40} {'ERROR':>10}")
    print()
    return wavs


def main():
    parser = argparse.ArgumentParser(description="G1 WAV Player (runs on robot)")
    parser.add_argument("path", nargs="?", help="WAV file or directory path")
    parser.add_argument("--list", metavar="DIR", help="List WAV files in directory")
    parser.add_argument("--all", metavar="DIR", help="Play all WAV files in directory")
    parser.add_argument("--volume", "-v", type=int, default=100, help="Volume 0-100")
    parser.add_argument("--interface", "-i", default="eth0", help="Network interface")
    parser.add_argument("--pause", "-p", type=float, default=1.0,
                        help="Pause between files in --all mode (seconds)")

    args = parser.parse_args()

    if args.list:
        list_wavs(args.list)
        return

    if args.all:
        wavs = sorted(Path(args.all).glob("*.wav"))
        if not wavs:
            print(f"No WAV files in {args.all}")
            return

        client = init_client(args.interface, args.volume)
        print(f"\nPlaying {len(wavs)} files from {args.all}\n")

        for i, wav_path in enumerate(wavs):
            pcm = load_wav(str(wav_path))
            play_wav(client, pcm, label=f"[{i+1}/{len(wavs)}] {wav_path.stem}")
            if i < len(wavs) - 1:
                time.sleep(args.pause)

        print(f"\nAll done.")
        return

    if not args.path:
        parser.print_help()
        return

    path = Path(args.path)
    if not path.exists():
        print(f"File not found: {path}")
        sys.exit(1)

    client = init_client(args.interface, args.volume)
    pcm = load_wav(str(path))
    play_wav(client, pcm, label=path.stem)


if __name__ == "__main__":
    main()