#!/usr/bin/env python3
"""
G1 VOICE RECORDER & REPLAYER
-----------------------------
Records voice from laptop microphone, saves as WAV, and replays
on the Unitree G1 robot via AudioClient.PlayStream().

Modes:
  record   - Record from mic and save to DataG1/<name>.wav
  replay   - Play a saved WAV file on G1 speaker
  list     - List all saved recordings
  live     - Record then immediately replay on G1

Audio Format:
  16-bit PCM, 16000 Hz, Mono (G1 voice service expects this)

Usage:
  python3 g1_voice_recorder.py enp3s0 record  --name greeting --seconds 5
  python3 g1_voice_recorder.py enp3s0 replay  --name greeting
  python3 g1_voice_recorder.py enp3s0 list
  python3 g1_voice_recorder.py enp3s0 live    --seconds 5
"""

import time
import sys
import wave
import argparse
import numpy as np
import os
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
G1_LOOTAH_DIR = SCRIPT_DIR.parent
DATA_DIR = SCRIPT_DIR / "DataG1"

# Add unitree_webrtc_connect to path
import sys
_webrtc_path = str(G1_LOOTAH_DIR / "unitree_webrtc_connect")
if _webrtc_path not in sys.path:
    sys.path.insert(0, _webrtc_path)

# G1 AudioClient PCM format
SAMPLE_RATE = 16000
CHANNELS = 1
SAMPLE_WIDTH = 2  # 16-bit
CHUNK_SIZE = 4096  # frames per buffer


def list_recordings():
    """List all saved WAV recordings."""
    wavs = sorted(DATA_DIR.glob("*.wav"))
    if not wavs:
        print("No recordings found in", DATA_DIR)
        return []

    print(f"\n{'Name':<30} {'Duration':>10} {'Size':>10}")
    print("-" * 52)
    for p in wavs:
        try:
            with wave.open(str(p), "rb") as wf:
                frames = wf.getnframes()
                rate = wf.getframerate()
                dur = frames / rate
            size_kb = p.stat().st_size / 1024
            print(f"{p.stem:<30} {dur:>8.1f}s {size_kb:>8.1f}KB")
        except Exception as e:
            print(f"{p.stem:<30} {'ERROR':>10} {str(e)}")
    print()
    return wavs


def _find_supported_rate(pa, device_index: int):
    """Find a sample rate the device actually supports."""
    import pyaudio
    info = pa.get_device_info_by_index(device_index)
    # Try common rates, prefer the device's default first
    candidates = [int(info["defaultSampleRate"]), 48000, 44100, 32000, 16000, 8000]
    for rate in candidates:
        try:
            supported = pa.is_format_supported(
                rate, input_device=device_index,
                input_channels=1, input_format=pyaudio.paInt16,
            )
            if supported:
                return rate
        except ValueError:
            continue
    # Fallback to device default
    return int(info["defaultSampleRate"])


def record_voice(name: str, seconds: float, device_index: int = None):
    """Record audio from microphone and save as WAV."""
    import pyaudio

    DATA_DIR.mkdir(parents=True, exist_ok=True)
    out_path = DATA_DIR / f"{name}.wav"

    pa = pyaudio.PyAudio()

    # Show available input devices — prefer 'default' or 'pipewire' over raw hw:
    if device_index is None:
        print("\nAvailable input devices:")
        best_idx = None
        for i in range(pa.get_device_count()):
            info = pa.get_device_info_by_index(i)
            if info["maxInputChannels"] > 0:
                dname = info["name"]
                marker = ""
                # Prefer 'default' or 'pipewire' — they handle mixing/resampling properly
                if best_idx is None:
                    if dname in ("default", "pipewire") or "pulse" in dname.lower():
                        best_idx = i
                        marker = " [SELECTED]"
                elif i == best_idx:
                    marker = " [SELECTED]"
                print(f"  [{i}] {dname} (in={info['maxInputChannels']}, "
                      f"default={int(info['defaultSampleRate'])}Hz){marker}")
        # If no preferred device found, fall back to first available
        if best_idx is None:
            for i in range(pa.get_device_count()):
                if pa.get_device_info_by_index(i)["maxInputChannels"] > 0:
                    best_idx = i
                    break
        device_index = best_idx
        print(f"\nUsing device [{device_index}] — {pa.get_device_info_by_index(device_index)['name']}")

    # Find a sample rate the hardware supports
    rec_rate = _find_supported_rate(pa, device_index)
    needs_resample = rec_rate != SAMPLE_RATE

    print(f"\nRecording: {seconds}s @ {rec_rate}Hz mono 16-bit"
          + (f" (will resample to {SAMPLE_RATE}Hz)" if needs_resample else ""))
    print(f"Output:    {out_path}")
    print("Press Ctrl+C to stop early.\n")

    stream = pa.open(
        format=pyaudio.paInt16,
        channels=CHANNELS,
        rate=rec_rate,
        input=True,
        input_device_index=device_index,
        frames_per_buffer=CHUNK_SIZE,
    )

    frames = []
    total_chunks = int(rec_rate / CHUNK_SIZE * seconds)
    start = time.time()

    try:
        for i in range(total_chunks):
            data = stream.read(CHUNK_SIZE, exception_on_overflow=False)
            frames.append(data)

            elapsed = time.time() - start
            bar_len = 30
            progress = min(elapsed / seconds, 1.0)
            filled = int(bar_len * progress)
            bar = "█" * filled + "░" * (bar_len - filled)
            print(f"\r  [{bar}] {elapsed:.1f}s / {seconds:.1f}s", end="", flush=True)

    except KeyboardInterrupt:
        print("\n  Recording stopped early.")

    print()

    stream.stop_stream()
    stream.close()
    pa.terminate()

    # Convert raw frames to numpy
    raw_audio = np.frombuffer(b"".join(frames), dtype=np.int16)

    # Resample to G1 target rate if needed
    if needs_resample:
        target_len = int(len(raw_audio) * SAMPLE_RATE / rec_rate)
        raw_audio = np.interp(
            np.linspace(0, len(raw_audio), target_len, endpoint=False),
            np.arange(len(raw_audio)),
            raw_audio.astype(np.float64),
        ).astype(np.int16)
        print(f"  Resampled: {rec_rate}Hz -> {SAMPLE_RATE}Hz")

    # Save WAV at G1 target rate
    wf = wave.open(str(out_path), "wb")
    wf.setnchannels(CHANNELS)
    wf.setsampwidth(SAMPLE_WIDTH)
    wf.setframerate(SAMPLE_RATE)
    wf.writeframes(raw_audio.tobytes())
    wf.close()

    duration = len(raw_audio) / SAMPLE_RATE
    size_kb = out_path.stat().st_size / 1024
    print(f"Saved: {out_path.name} ({duration:.1f}s, {size_kb:.1f}KB)")
    return out_path


def replay_on_g1(name: str, robot_ip: str, user: str = "unitree"):
    """Replay a saved WAV file on the G1 built-in speaker.

    SCP the WAV + a tiny player script to the robot, then SSH and run it
    using the robot's SDK which has AudioClient.PlayStream().
    """
    import subprocess

    wav_path = DATA_DIR / f"{name}.wav"
    if not wav_path.exists():
        print(f"Recording not found: {wav_path}")
        print("Available recordings:")
        list_recordings()
        sys.exit(1)

    # Read WAV info
    with wave.open(str(wav_path), "rb") as wf:
        framerate = wf.getframerate()
        n_channels = wf.getnchannels()
        sampwidth = wf.getsampwidth()
        n_frames = wf.getnframes()

    duration = n_frames / framerate
    print(f"\nLoaded: {wav_path.name}")
    print(f"  Format:   {framerate}Hz, {n_channels}ch, {sampwidth * 8}-bit")
    print(f"  Duration: {duration:.1f}s")

    remote_wav = f"/tmp/{wav_path.name}"
    remote_script = "/tmp/_g1_play_wav.py"
    target = f"{user}@{robot_ip}"
    ssh_opts = ["-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=5"]

    # Player script that runs ON the robot — uses _CallRequestWithParamAndBin
    # (single call, unique stream_id, format params in JSON). PlayStream is
    # unreliable — see voice_note.txt.
    player_script = '''\
import sys, time, wave, json, numpy as np
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  # G1 speaker expects 16kHz mono 16-bit

wav_path = sys.argv[1]
ChannelFactoryInitialize(0, "eth0")
client = AudioClient()
client.SetTimeout(10.0)
client.Init()
client.SetVolume(100)

with wave.open(wav_path, "rb") as wf:
    rate = wf.getframerate()
    nch = wf.getnchannels()
    pcm_raw = wf.readframes(wf.getnframes())

audio = np.frombuffer(pcm_raw, 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)

pcm_bytes = audio.tobytes()
duration = len(audio) / TARGET_RATE

# Reset any previous stream, then send ALL audio in ONE call with a unique sid.
client._Call(ROBOT_API_ID_AUDIO_STOP_PLAY, json.dumps({"app_name": "voice_rec"}))
time.sleep(0.3)

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

print(f"Playing {duration:.1f}s via _CallRequestWithParamAndBin (16kHz mono)...")
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": "voice_rec"}))
print(f"Done. ({time.time() - start:.1f}s)")
'''

    # Step 1: Copy WAV + player script to robot
    print(f"\nCopying to {target}...")

    # Write player script to local temp file
    import tempfile
    local_script = os.path.join(tempfile.gettempdir(), "_g1_play_wav.py")
    with open(local_script, "w") as f:
        f.write(player_script)

    # SCP WAV file
    ret = subprocess.run(
        ["scp"] + ssh_opts + [str(wav_path), f"{target}:{remote_wav}"],
        capture_output=True, text=True,
    )
    if ret.returncode != 0:
        print(f"  SCP WAV failed: {ret.stderr.strip()}")
        os.unlink(local_script)
        sys.exit(1)

    # SCP player script
    ret = subprocess.run(
        ["scp"] + ssh_opts + [local_script, f"{target}:{remote_script}"],
        capture_output=True, text=True,
    )
    os.unlink(local_script)

    if ret.returncode != 0:
        print(f"  SCP script failed: {ret.stderr.strip()}")
        sys.exit(1)
    print("  Copied.")

    # Step 2: Run player script on robot via SSH
    # Use full path to gemini env python (conda run fails in non-interactive SSH)
    play_cmd = (
        f"/home/unitree/miniconda3/envs/gemini/bin/python3 "
        f"{remote_script} {remote_wav}"
    )

    print(f"Playing on G1... ({duration:.1f}s)")
    proc = subprocess.Popen(
        ["ssh"] + ssh_opts + [target, play_cmd],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
    )

    start = time.time()
    while proc.poll() is None:
        elapsed = time.time() - start
        progress = min(elapsed / duration, 1.0)
        bar_len = 30
        filled = int(bar_len * progress)
        bar = "█" * filled + "░" * (bar_len - filled)
        print(f"\r  [{bar}] {min(elapsed, duration):.1f}s / {duration:.1f}s",
              end="", flush=True)
        time.sleep(0.2)

    output = proc.stdout.read().decode().strip()
    if output:
        print(f"\n  Robot: {output}")

    if proc.returncode != 0:
        print(f"\n  SSH exit code: {proc.returncode}")
    else:
        print(f"\nPlayback complete.")

    # Cleanup
    subprocess.run(
        ["ssh"] + ssh_opts + [target, f"rm -f {remote_wav} {remote_script}"],
        capture_output=True,
    )


def live_record_and_replay(seconds: float, robot_ip: str, device_index: int = None, user: str = "unitree"):
    """Record voice then immediately replay on G1."""
    name = f"live_{int(time.time())}"
    print("=== STEP 1: RECORD ===")
    record_voice(name, seconds, device_index)

    print("\n=== STEP 2: REPLAY ON G1 ===")
    replay_on_g1(name, robot_ip, user)


def play_local(name: str):
    """Play a saved WAV file on the local workstation speaker."""
    import pyaudio

    wav_path = DATA_DIR / f"{name}.wav"
    if not wav_path.exists():
        print(f"Recording not found: {wav_path}")
        list_recordings()
        sys.exit(1)

    with wave.open(str(wav_path), "rb") as wf:
        n_channels = wf.getnchannels()
        sampwidth = wf.getsampwidth()
        framerate = wf.getframerate()
        n_frames = wf.getnframes()

        duration = n_frames / framerate
        print(f"\nPlaying locally: {wav_path.name}")
        print(f"  Format:   {framerate}Hz, {n_channels}ch, {sampwidth * 8}-bit")
        print(f"  Duration: {duration:.1f}s\n")

        pa = pyaudio.PyAudio()
        stream = pa.open(
            format=pa.get_format_from_width(sampwidth),
            channels=n_channels,
            rate=framerate,
            output=True,
        )

        start = time.time()
        chunk = wf.readframes(CHUNK_SIZE)
        while chunk:
            stream.write(chunk)
            elapsed = time.time() - start
            progress = min(elapsed / duration, 1.0)
            bar_len = 30
            filled = int(bar_len * progress)
            bar = "█" * filled + "░" * (bar_len - filled)
            print(f"\r  [{bar}] {elapsed:.1f}s / {duration:.1f}s", end="", flush=True)
            chunk = wf.readframes(CHUNK_SIZE)

        print(f"\n  Done.")
        stream.stop_stream()
        stream.close()
        pa.terminate()


def main():
    parser = argparse.ArgumentParser(
        description="G1 Voice Recorder & Replayer",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s record    --name hello --seconds 5
  %(prog)s playlocal --name hello
  %(prog)s replay    --name hello --ip 192.168.123.164
  %(prog)s list
  %(prog)s live      --seconds 5 --ip 192.168.123.164
        """,
    )
    parser.add_argument("mode", choices=["record", "replay", "list", "live", "playlocal"],
                        help="Operation mode")
    parser.add_argument("--name", "-n", default="recording",
                        help="Recording name (default: recording)")
    parser.add_argument("--seconds", "-s", type=float, default=5.0,
                        help="Recording duration in seconds (default: 5)")
    parser.add_argument("--device", "-d", type=int, default=None,
                        help="Audio input device index (default: auto-detect)")
    parser.add_argument("--ip", default="192.168.123.164",
                        help="G1 robot IP for replay/live (default: 192.168.123.164)")
    parser.add_argument("--user", "-u", default="unitree",
                        help="SSH user on G1 (default: unitree)")

    args = parser.parse_args()

    if args.mode == "list":
        list_recordings()

    elif args.mode == "record":
        record_voice(args.name, args.seconds, args.device)

    elif args.mode == "playlocal":
        play_local(args.name)

    elif args.mode == "replay":
        replay_on_g1(args.name, args.ip, args.user)

    elif args.mode == "live":
        live_record_and_replay(args.seconds, args.ip, args.device, args.user)


if __name__ == "__main__":
    main()
