#!/usr/bin/env python3
"""
G1 Voice Deploy — Record on workstation, deploy & play on G1 robot.
Also sets up mic on robot and manages remote recordings.

Usage (run from workstation):
  python3 g1_voice_deploy.py setup-mic          # Unmute & configure wireless mic on robot
  python3 g1_voice_deploy.py record --name hi    # Record from workstation mic, save locally
  python3 g1_voice_deploy.py record-robot --name arabic_greeting --seconds 5  # Record from robot mic
  python3 g1_voice_deploy.py play --name hi      # Deploy & play on G1 speaker
  python3 g1_voice_deploy.py play-robot --path "/home/unitree/SanadVoice/recorded voices/welcome_single.wav"
  python3 g1_voice_deploy.py tts --text "Hello"  # TTS on G1 speaker
  python3 g1_voice_deploy.py list                # List local recordings
  python3 g1_voice_deploy.py list-robot          # List recordings on robot
"""

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

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

ROBOT_IP = "192.168.123.164"
ROBOT_USER = "unitree"
ROBOT_PYTHON = "/home/unitree/miniconda3/envs/gemini/bin/python3"
ROBOT_VOICES_DIR = "/home/unitree/SanadVoice/recorded voices"
ROBOT_DDS_IFACE = "eth0"

SAMPLE_RATE = 16000
CHANNELS = 1
SAMPLE_WIDTH = 2
CHUNK_SIZE = 4096

SSH_OPTS = ["-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=5"]


def ssh_cmd(cmd: str, capture=True):
    """Run a command on the robot via SSH."""
    target = f"{ROBOT_USER}@{ROBOT_IP}"
    result = subprocess.run(
        ["ssh"] + SSH_OPTS + [target, cmd],
        capture_output=capture, text=True,
    )
    if capture:
        return result.returncode, result.stdout.strip(), result.stderr.strip()
    return result.returncode, "", ""


def scp_to_robot(local_path: str, remote_path: str):
    """Copy a file to the robot."""
    target = f"{ROBOT_USER}@{ROBOT_IP}"
    ret = subprocess.run(
        ["scp"] + SSH_OPTS + [local_path, f"{target}:{remote_path}"],
        capture_output=True, text=True,
    )
    return ret.returncode == 0


def scp_from_robot(remote_path: str, local_path: str):
    """Copy a file from the robot."""
    target = f"{ROBOT_USER}@{ROBOT_IP}"
    ret = subprocess.run(
        ["scp"] + SSH_OPTS + [f"{target}:{remote_path}", local_path],
        capture_output=True, text=True,
    )
    return ret.returncode == 0


# ─── SETUP MIC ──────────────────────────────────────────────

def cmd_setup_mic(args):
    """Unmute and configure the wireless mic on the robot."""
    print("Setting up wireless mic on G1...\n")

    # List sources
    code, out, _ = ssh_cmd("pactl list sources short")
    print("PulseAudio sources:")
    print(out)
    print()

    # Find wireless mic source
    mic_source = None
    for line in out.split("\n"):
        if "Wireless" in line or "Hollyland" in line:
            mic_source = line.split("\t")[1] if "\t" in line else None
            mic_index = line.split("\t")[0] if "\t" in line else None
            break

    if not mic_source:
        print("Wireless mic not found. Is it plugged in?")
        return

    print(f"Found mic: {mic_source} (index {mic_index})")

    # Set as default, unmute, set volume
    ssh_cmd(f"pactl set-default-source {mic_source}")
    print("  Set as default source.")

    ssh_cmd(f"pactl set-source-mute {mic_index} 0")
    print("  Unmuted.")

    ssh_cmd(f"pactl set-source-volume {mic_index} 100%")
    print("  Volume set to 100%.")

    # Check mute status
    code, out, _ = ssh_cmd(f'pactl list sources | grep -A 3 "Wireless" | grep Mute')
    print(f"  Status: {out.strip()}")

    # Verify recording
    print("\n  Verifying mic (2 second test)...")
    ssh_cmd("timeout 2 parec -d 3 --format=s16le --rate=16000 --channels=1 --raw > /tmp/_mic_verify.pcm")
    code, out, _ = ssh_cmd(
        f'{ROBOT_PYTHON} -c "'
        'import numpy as np;'
        'a = np.fromfile(\'/tmp/_mic_verify.pcm\', dtype=np.int16);'
        'print(f\'samples={len(a)} std={a.std():.0f}\')'
        '"'
    )
    print(f"  Result: {out}")
    if "std=0" in out:
        print("  WARNING: Mic is silent! Check transmitter is powered on.")
    else:
        print("  Mic is working!")


# ─── RECORD ON WORKSTATION ───────────────────────────────────

def cmd_record(args):
    """Record from workstation microphone."""
    import pyaudio

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

    pa = pyaudio.PyAudio()

    # Find best input device
    device_index = args.device
    if device_index is None:
        for i in range(pa.get_device_count()):
            info = pa.get_device_info_by_index(i)
            if info["maxInputChannels"] > 0:
                name = info["name"]
                if name in ("default", "pipewire") or "pulse" in name.lower():
                    device_index = i
                    break
        if device_index is None:
            for i in range(pa.get_device_count()):
                if pa.get_device_info_by_index(i)["maxInputChannels"] > 0:
                    device_index = i
                    break

    # Find supported rate
    info = pa.get_device_info_by_index(device_index)
    rec_rate = int(info["defaultSampleRate"])
    print(f"Using device [{device_index}] {info['name']} @ {rec_rate}Hz")
    print(f"Recording {args.seconds}s... speak now!\n")

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

    frames = []
    total = int(rec_rate / CHUNK_SIZE * args.seconds)
    start = time.time()
    for i in range(total):
        frames.append(stream.read(CHUNK_SIZE, exception_on_overflow=False))
        elapsed = time.time() - start
        pct = min(elapsed / args.seconds, 1.0)
        bar = "█" * int(30 * pct) + "░" * (30 - int(30 * pct))
        print(f"\r  [{bar}] {elapsed:.1f}s / {args.seconds:.1f}s", end="", flush=True)

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

    # Resample to 16kHz
    audio = np.frombuffer(b"".join(frames), dtype=np.int16)
    if rec_rate != SAMPLE_RATE:
        tl = int(len(audio) * SAMPLE_RATE / rec_rate)
        audio = np.interp(
            np.linspace(0, len(audio), tl, endpoint=False),
            np.arange(len(audio)), audio.astype(np.float64),
        ).astype(np.int16)

    wf = wave.open(str(out_path), "wb")
    wf.setnchannels(1)
    wf.setsampwidth(2)
    wf.setframerate(SAMPLE_RATE)
    wf.writeframes(audio.tobytes())
    wf.close()

    print(f"Saved: {out_path} ({len(audio)/SAMPLE_RATE:.1f}s, {out_path.stat().st_size/1024:.1f}KB)")


# ─── RECORD ON ROBOT ─────────────────────────────────────────

def cmd_record_robot(args):
    """Record from the wireless mic on the robot."""
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    seconds = args.seconds
    remote_wav = f"/tmp/_rec_{args.name}.wav"
    local_path = DATA_DIR / f"{args.name}.wav"

    record_code = (
        f"import time, subprocess, wave, numpy as np;"
        f"print('Recording {seconds}s... speak now!');"
        f"proc = subprocess.Popen(['parec','-d','3','--format=s16le','--rate=16000','--channels=1','--raw'], stdout=subprocess.PIPE);"
        f"time.sleep({seconds});"
        f"proc.terminate();"
        f"raw = proc.stdout.read();"
        f"audio = np.frombuffer(raw, dtype=np.int16);"
        f"print(f'Recorded {{len(audio)}} samples, std={{audio.std():.0f}}');"
        f"wf = wave.open('{remote_wav}', 'wb');"
        f"wf.setnchannels(1); wf.setsampwidth(2); wf.setframerate(16000);"
        f"wf.writeframes(audio.tobytes()); wf.close();"
        f"print('Saved on robot')"
    )

    print(f"Recording {seconds}s from robot wireless mic...")
    code, out, err = ssh_cmd(f'{ROBOT_PYTHON} -c "{record_code}"')
    print(f"  {out}")
    if err:
        print(f"  stderr: {err}")

    # Download to workstation
    print(f"Downloading to {local_path}...")
    if scp_from_robot(remote_wav, str(local_path)):
        print(f"  Saved: {local_path}")
        ssh_cmd(f"rm -f {remote_wav}")
    else:
        print("  Download failed!")

    # Also save on robot in recorded voices dir
    remote_final = f"{ROBOT_VOICES_DIR}/{args.name}.wav"
    ssh_cmd(f'cp "{remote_wav}" "{remote_final}" 2>/dev/null')


# ─── PLAY ON G1 ──────────────────────────────────────────────

def cmd_play(args):
    """Deploy a local WAV file to G1 and play on built-in speaker."""
    wav_path = DATA_DIR / f"{args.name}.wav"
    if not wav_path.exists():
        print(f"Not found: {wav_path}")
        cmd_list(args)
        return

    with wave.open(str(wav_path), "rb") as wf:
        dur = wf.getnframes() / wf.getframerate()
    print(f"Loaded: {wav_path.name} ({dur:.1f}s)")

    # SCP to robot
    remote_wav = f"/tmp/_play_{args.name}.wav"
    print(f"Copying to robot...")
    if not scp_to_robot(str(wav_path), remote_wav):
        print("  SCP failed!")
        return
    print("  Copied.")

    # Play on robot
    _play_remote_wav(remote_wav, dur)

    # Cleanup
    ssh_cmd(f"rm -f {remote_wav}")


def cmd_play_robot(args):
    """Play a WAV file that already exists on the robot."""
    # Get duration
    code, out, _ = ssh_cmd(
        f'{ROBOT_PYTHON} -c "import wave; '
        f"wf = wave.open('{args.path}', 'rb'); "
        f"print(wf.getnframes() / wf.getframerate()); wf.close()"
        '"'
    )
    try:
        dur = float(out)
    except ValueError:
        print(f"Cannot read: {args.path}")
        print(f"  {out}")
        return

    print(f"Playing: {args.path} ({dur:.1f}s)")
    _play_remote_wav(args.path, dur)


def _play_remote_wav(remote_wav: str, duration: float):
    """Play a WAV file on the G1 speaker (file must already be on robot)."""
    play_code = (
        "import 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 *;"
        f"ChannelFactoryInitialize(0, '{ROBOT_DDS_IFACE}');"
        "c = AudioClient(); c.SetTimeout(10.0); c.Init(); c.SetVolume(100);"
        "c._Call(ROBOT_API_ID_AUDIO_STOP_PLAY, json.dumps({'app_name':'p'}));"
        "time.sleep(0.3);"
        f"wf = wave.open('{remote_wav}', 'rb');"
        "rate = wf.getframerate(); nch = wf.getnchannels();"
        "audio = np.frombuffer(wf.readframes(wf.getnframes()), dtype=np.int16); wf.close();"
        "audio = audio.reshape(-1, 2).mean(axis=1).astype(np.int16) if nch == 2 else audio;"
        "tl = int(len(audio) * 16000 / rate) if rate != 16000 else len(audio);"
        "audio = np.interp(np.linspace(0, len(audio), tl, endpoint=False), np.arange(len(audio)), audio.astype(np.float64)).astype(np.int16) if rate != 16000 else audio;"
        "pcm = audio.tobytes();"
        "sid = f's_{int(time.time()*1000)}';"
        "param = json.dumps({'app_name':'p','stream_id':sid,'sample_rate':16000,'channels':1,'bits_per_sample':16});"
        "c._CallRequestWithParamAndBin(ROBOT_API_ID_AUDIO_START_PLAY, param, list(pcm));"
        f"time.sleep({duration + 1});"
        "c._Call(ROBOT_API_ID_AUDIO_STOP_PLAY, json.dumps({'app_name':'p'}));"
        "print('played')"
    )

    print(f"Playing on G1... ({duration:.1f}s)")
    proc = subprocess.Popen(
        ["ssh"] + SSH_OPTS + [f"{ROBOT_USER}@{ROBOT_IP}",
         f'{ROBOT_PYTHON} -c "{play_code}"'],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
    )

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

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


# ─── TTS ──────────────────────────────────────────────────────

def cmd_tts(args):
    """Play text-to-speech on G1 speaker (English/Chinese only)."""
    tts_code = (
        "import json;"
        "from unitree_sdk2py.core.channel import ChannelFactoryInitialize;"
        "from unitree_sdk2py.g1.audio.g1_audio_client import AudioClient;"
        f"ChannelFactoryInitialize(0, '{ROBOT_DDS_IFACE}');"
        "c = AudioClient(); c.SetTimeout(10.0); c.Init();"
        f"c.SetVolume({args.volume});"
        f"c.TtsMaker('{args.text}', {args.speaker_id});"
        "print('ok')"
    )
    print(f"TTS: \"{args.text}\" (speaker_id={args.speaker_id}, vol={args.volume})")
    code, out, _ = ssh_cmd(f'{ROBOT_PYTHON} -c "{tts_code}"')
    if "ok" in out:
        print("  Sent.")
    else:
        print(f"  Result: {out}")


# ─── LIST ─────────────────────────────────────────────────────

def cmd_list(args):
    """List local recordings."""
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    wavs = sorted(DATA_DIR.glob("*.wav"))
    if not wavs:
        print(f"No recordings in {DATA_DIR}")
        return
    print(f"\n{'#':<4} {'Name':<35} {'Duration':>8} {'Size':>10}")
    print("-" * 60)
    for i, p in enumerate(wavs):
        try:
            with wave.open(str(p), "rb") as wf:
                dur = wf.getnframes() / wf.getframerate()
            size = p.stat().st_size / 1024
            print(f"{i:<4} {p.stem:<35} {dur:>6.1f}s {size:>8.1f}KB")
        except Exception:
            print(f"{i:<4} {p.stem:<35} {'ERROR':>8}")
    print()


def cmd_list_robot(args):
    """List recordings on the robot."""
    code, out, _ = ssh_cmd(
        f'{ROBOT_PYTHON} -c "'
        "import wave, os; from pathlib import Path;"
        f"d = Path('{ROBOT_VOICES_DIR}');"
        "wavs = sorted(d.glob('*.wav'));"
        "print(f'{{len(wavs)}} files in {d}');"
        "for p in wavs:"
        "    try:"
        "        wf = wave.open(str(p),'rb'); dur = wf.getnframes()/wf.getframerate(); wf.close();"
        "        print(f'{p.stem:<45} {dur:>6.1f}s')"
        "    except: print(f'{p.stem:<45} ERROR')"
        '"'
    )
    print(out)


# ─── MAIN ─────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="G1 Voice Deploy — Record, deploy & play audio on G1 robot",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s setup-mic
  %(prog)s record --name greeting --seconds 5
  %(prog)s record-robot --name arabic_hi --seconds 5
  %(prog)s play --name greeting
  %(prog)s play-robot --path "/home/unitree/SanadVoice/recorded voices/welcome_single.wav"
  %(prog)s tts --text "Hello, welcome!"
  %(prog)s list
  %(prog)s list-robot
        """,
    )

    sub = parser.add_subparsers(dest="command")

    # setup-mic
    sub.add_parser("setup-mic", help="Unmute & configure wireless mic on robot")

    # record
    p = sub.add_parser("record", help="Record from workstation mic")
    p.add_argument("--name", "-n", required=True, help="Recording name")
    p.add_argument("--seconds", "-s", type=float, default=5.0)
    p.add_argument("--device", "-d", type=int, default=None)

    # record-robot
    p = sub.add_parser("record-robot", help="Record from robot wireless mic")
    p.add_argument("--name", "-n", required=True, help="Recording name")
    p.add_argument("--seconds", "-s", type=float, default=5.0)

    # play
    p = sub.add_parser("play", help="Deploy & play local WAV on G1 speaker")
    p.add_argument("--name", "-n", required=True, help="Recording name")

    # play-robot
    p = sub.add_parser("play-robot", help="Play a WAV file already on robot")
    p.add_argument("--path", "-p", required=True, help="Full path on robot")

    # tts
    p = sub.add_parser("tts", help="Text-to-speech on G1 (English/Chinese)")
    p.add_argument("--text", "-t", required=True, help="Text to speak")
    p.add_argument("--speaker-id", type=int, default=0, help="Voice ID (0,1,2)")
    p.add_argument("--volume", "-v", type=int, default=100)

    # list
    sub.add_parser("list", help="List local recordings")
    sub.add_parser("list-robot", help="List recordings on robot")

    args = parser.parse_args()

    commands = {
        "setup-mic": cmd_setup_mic,
        "record": cmd_record,
        "record-robot": cmd_record_robot,
        "play": cmd_play,
        "play-robot": cmd_play_robot,
        "tts": cmd_tts,
        "list": cmd_list,
        "list-robot": cmd_list_robot,
    }

    if args.command in commands:
        commands[args.command](args)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()