#!/usr/bin/env python3
"""
G1 EDGE-TTS — Text to G1 speaker via Microsoft Edge TTS (no API key).

Pipeline matches Marcus's working path documented in voice_note.txt:
    edge-tts  →  MP3 bytes  →  pydub 16 kHz mono WAV
              →  _CallRequestWithParamAndBin on G1 AudioClient

Voices (defaults):
    Arabic   ar-AE-HamdanNeural   (UAE male)
    English  en-US-GuyNeural      (US male)

Runs on robot (has AudioClient) OR on workstation (save WAV only with --save).

Requirements:
    pip install edge-tts pydub numpy
    (on robot use the gemini conda env which has the audio SDK methods)

Usage:
    # Speak directly on G1 speaker:
    python3 g1_edge_tts.py "مرحبا بكم" --lang ar
    python3 g1_edge_tts.py "Hello welcome to Lootah" --lang en

    # Just save a WAV (no robot needed):
    python3 g1_edge_tts.py "Good morning" --save greeting.wav

    # Override voice:
    python3 g1_edge_tts.py "test" --lang ar --voice ar-SA-HamedNeural
"""

from __future__ import annotations

import argparse
import asyncio
import io
import json
import sys
import time
import wave
from pathlib import Path

import numpy as np

DEFAULT_VOICES = {
    "ar": "ar-AE-HamdanNeural",
    "en": "en-US-GuyNeural",
}
TARGET_RATE = 16000
APP_NAME = "edge_tts"


async def _synthesize(text: str, voice: str) -> bytes:
    """Call edge-tts and return raw MP3 bytes."""
    import edge_tts  # lazy import — only needed when actually speaking

    communicate = edge_tts.Communicate(text=text, voice=voice)
    buf = io.BytesIO()
    async for chunk in communicate.stream():
        if chunk["type"] == "audio":
            buf.write(chunk["data"])
    return buf.getvalue()


def _mp3_to_pcm16k(mp3: bytes) -> np.ndarray:
    """Decode MP3 → 16 kHz mono int16."""
    from pydub import AudioSegment  # lazy — pulls ffmpeg at import time

    seg = AudioSegment.from_file(io.BytesIO(mp3), format="mp3")
    seg = seg.set_channels(1).set_frame_rate(TARGET_RATE).set_sample_width(2)
    return np.frombuffer(seg.raw_data, dtype=np.int16)


def save_wav(pcm: np.ndarray, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with wave.open(str(path), "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(TARGET_RATE)
        wf.writeframes(pcm.tobytes())


def _play_on_g1(pcm: np.ndarray, interface: str = "eth0", volume: int = 100) -> float:
    """Single-call _CallRequestWithParamAndBin — the working method."""
    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, interface)
    c = AudioClient()
    c.SetTimeout(10.0)
    c.Init()
    c.SetVolume(volume)

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

    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()
    c._CallRequestWithParamAndBin(
        ROBOT_API_ID_AUDIO_START_PLAY, param, list(pcm_bytes)
    )
    time.sleep(duration + 0.5)
    c._Call(ROBOT_API_ID_AUDIO_STOP_PLAY, json.dumps({"app_name": APP_NAME}))
    return time.time() - start


def main() -> int:
    p = argparse.ArgumentParser(description="Edge-TTS → G1 speaker (AR + EN)")
    p.add_argument("text", help="Text to speak (wrap in quotes)")
    p.add_argument("--lang", choices=["ar", "en"], default="en",
                   help="Language (default: en)")
    p.add_argument("--voice", default=None,
                   help="Override voice (e.g. ar-SA-HamedNeural)")
    p.add_argument("--save", metavar="PATH",
                   help="Also save the synthesized 16 kHz WAV to PATH")
    p.add_argument("--no-play", action="store_true",
                   help="Skip G1 playback; useful with --save from workstation")
    p.add_argument("--interface", default="eth0", help="DDS interface (default: eth0)")
    p.add_argument("--volume", type=int, default=100, help="Speaker volume 0-100")
    args = p.parse_args()

    voice = args.voice or DEFAULT_VOICES[args.lang]
    print(f"Synthesizing ({args.lang}, voice={voice}) ...")
    t0 = time.time()
    try:
        mp3 = asyncio.run(_synthesize(args.text, voice))
    except Exception as exc:
        print(f"edge-tts failed: {exc}", file=sys.stderr)
        return 1
    pcm = _mp3_to_pcm16k(mp3)
    dur = len(pcm) / TARGET_RATE
    print(f"  {len(mp3):,} MP3 bytes -> {dur:.1f}s WAV ({time.time() - t0:.1f}s synth)")

    if args.save:
        save_path = Path(args.save)
        save_wav(pcm, save_path)
        print(f"  saved {save_path} ({save_path.stat().st_size/1024:.1f} KB)")

    if args.no_play:
        return 0

    try:
        elapsed = _play_on_g1(pcm, args.interface, args.volume)
        print(f"Played on G1 ({elapsed:.1f}s).")
    except Exception as exc:
        print(f"G1 playback failed: {exc}", file=sys.stderr)
        print("Hint: run this on the robot inside the gemini conda env.")
        return 2
    return 0


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