#!/usr/bin/env python3
"""
G1 Arabic TTS — Offline text-to-speech with playback on G1 built-in speaker.
Supports Piper TTS (recommended) and tts_arabic as backends.

Install on robot (gemini env):
  pip install piper-tts
  # First run auto-downloads Arabic voice model (~50MB)

Usage (run ON the robot):
  python3 g1_tts_arabic.py "مرحبا بكم في لوتاه"
  python3 g1_tts_arabic.py "مرحبا" --save "/home/unitree/SanadVoice/recorded voices/marhaba.wav"
  python3 g1_tts_arabic.py --interactive
  python3 g1_tts_arabic.py "Hello welcome" --lang en
  python3 g1_tts_arabic.py --list-voices
"""

import argparse
import io
import json
import os
import sys
import time
import wave
import numpy as np

TARGET_RATE = 16000
DDS_IFACE = "eth0"


def _init_g1_audio():
    """Initialize G1 AudioClient."""
    from unitree_sdk2py.core.channel import ChannelFactoryInitialize
    from unitree_sdk2py.g1.audio.g1_audio_client import AudioClient
    ChannelFactoryInitialize(0, DDS_IFACE)
    client = AudioClient()
    client.SetTimeout(10.0)
    client.Init()
    client.SetVolume(100)
    return client


def _play_on_g1(client, audio_16k: np.ndarray):
    """Play 16kHz mono int16 audio on G1 speaker."""
    from unitree_sdk2py.g1.audio.g1_audio_api import (
        ROBOT_API_ID_AUDIO_START_PLAY,
        ROBOT_API_ID_AUDIO_STOP_PLAY,
    )

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

    pcm = audio_16k.tobytes()
    sid = f"s_{int(time.time() * 1000)}"
    param = json.dumps({
        "app_name": "tts",
        "stream_id": sid,
        "sample_rate": TARGET_RATE,
        "channels": 1,
        "bits_per_sample": 16,
    })
    client._CallRequestWithParamAndBin(ROBOT_API_ID_AUDIO_START_PLAY, param, list(pcm))

    duration = len(audio_16k) / TARGET_RATE
    time.sleep(duration + 0.5)
    client._Call(ROBOT_API_ID_AUDIO_STOP_PLAY, json.dumps({"app_name": "tts"}))
    return duration


def _resample(audio: np.ndarray, src_rate: int) -> np.ndarray:
    """Resample to 16kHz."""
    if src_rate == TARGET_RATE:
        return audio
    tl = int(len(audio) * TARGET_RATE / src_rate)
    return np.interp(
        np.linspace(0, len(audio), tl, endpoint=False),
        np.arange(len(audio)),
        audio.astype(np.float64),
    ).astype(np.int16)


def _save_wav(path: str, audio: np.ndarray, rate: int = TARGET_RATE):
    """Save audio as WAV file."""
    wf = wave.open(path, "wb")
    wf.setnchannels(1)
    wf.setsampwidth(2)
    wf.setframerate(rate)
    wf.writeframes(audio.tobytes())
    wf.close()


# ─── PIPER TTS BACKEND ──────────────────────────────────────

def tts_piper(text: str, voice: str = "ar_JO-kareem-medium", speaker: int = None):
    """Generate speech using Piper TTS. Returns (audio_int16, sample_rate)."""
    try:
        import subprocess
        # Use piper CLI — most reliable on ARM64
        cmd = ["piper", "--model", voice, "--output_raw"]
        if speaker is not None:
            cmd += ["--speaker", str(speaker)]

        proc = subprocess.run(
            cmd, input=text.encode("utf-8"),
            capture_output=True, timeout=30,
        )

        if proc.returncode != 0:
            stderr = proc.stderr.decode()
            # If model not found, try downloading
            if "not found" in stderr.lower() or "No such file" in stderr:
                print(f"  Downloading voice model '{voice}'...")
                proc = subprocess.run(
                    cmd, input=text.encode("utf-8"),
                    capture_output=True, timeout=120,
                )

        if proc.returncode != 0:
            raise RuntimeError(f"Piper error: {proc.stderr.decode()[:200]}")

        # Piper outputs raw 16-bit mono PCM at 22050Hz by default
        audio = np.frombuffer(proc.stdout, dtype=np.int16)
        return audio, 22050

    except FileNotFoundError:
        raise RuntimeError("Piper not installed. Run: pip install piper-tts")


def tts_piper_python(text: str, voice: str = "ar_JO-kareem-medium"):
    """Generate speech using Piper Python API."""
    try:
        from piper import PiperVoice

        model_path = os.path.expanduser(f"~/.local/share/piper-voices/{voice}.onnx")
        if not os.path.exists(model_path):
            # Try auto-download via CLI first
            print(f"  Model not found at {model_path}")
            print(f"  Run: piper --model {voice} --output_raw <<< 'test'")
            raise FileNotFoundError(model_path)

        pv = PiperVoice.load(model_path)
        buf = io.BytesIO()
        with wave.open(buf, "wb") as wf:
            pv.synthesize(text, wf)

        buf.seek(0)
        with wave.open(buf, "rb") as wf:
            rate = wf.getframerate()
            audio = np.frombuffer(wf.readframes(wf.getnframes()), dtype=np.int16)

        return audio, rate

    except ImportError:
        raise RuntimeError("Piper not installed. Run: pip install piper-tts")


# ─── ESPEAK BACKEND (FALLBACK) ──────────────────────────────

def tts_espeak(text: str, lang: str = "ar", speed: int = 130):
    """Generate speech using espeak-ng. Low quality but always available."""
    import subprocess

    wav_path = "/tmp/_espeak_tts.wav"
    cmd = ["espeak-ng", "-v", lang, "-s", str(speed), "-w", wav_path, text]
    proc = subprocess.run(cmd, capture_output=True)

    if proc.returncode != 0:
        raise RuntimeError(f"espeak-ng error: {proc.stderr.decode()[:200]}")

    with wave.open(wav_path, "rb") as wf:
        rate = wf.getframerate()
        audio = np.frombuffer(wf.readframes(wf.getnframes()), dtype=np.int16)

    os.unlink(wav_path)
    return audio, rate


# ─── UNIFIED TTS ─────────────────────────────────────────────

def synthesize(text: str, backend: str = "auto", lang: str = "ar", voice: str = None):
    """Synthesize text to audio. Returns (audio_int16_16kHz, duration)."""

    if backend == "auto":
        # Try piper first, fall back to espeak
        for try_backend in ["piper", "espeak"]:
            try:
                return synthesize(text, backend=try_backend, lang=lang, voice=voice)
            except RuntimeError as e:
                print(f"  {try_backend} failed: {e}")
                continue
        raise RuntimeError("No TTS backend available!")

    if backend == "piper":
        if voice is None:
            voice = "ar_JO-kareem-medium" if lang == "ar" else "en_US-lessac-medium"
        audio, rate = tts_piper(text, voice=voice)

    elif backend == "espeak":
        audio, rate = tts_espeak(text, lang=lang)

    else:
        raise ValueError(f"Unknown backend: {backend}")

    # Resample to 16kHz
    audio_16k = _resample(audio, rate)
    duration = len(audio_16k) / TARGET_RATE

    return audio_16k, duration


# ─── COMMANDS ─────────────────────────────────────────────────

def cmd_speak(args):
    """Synthesize and play on G1."""
    text = args.text
    print(f'Text: "{text}"')
    print(f"Backend: {args.backend}, lang: {args.lang}")

    audio, duration = synthesize(text, backend=args.backend, lang=args.lang, voice=args.voice)
    print(f"Generated: {duration:.1f}s")

    if args.save:
        _save_wav(args.save, audio)
        print(f"Saved: {args.save}")

    if not args.no_play:
        print("Playing on G1...")
        client = _init_g1_audio()
        _play_on_g1(client, audio)
        print("Done.")


def cmd_interactive(args):
    """Interactive Arabic TTS loop."""
    print("G1 Arabic TTS — Interactive Mode")
    print("Type Arabic (or English) text, press Enter to speak.")
    print("Commands: /lang ar|en, /voice <name>, /backend piper|espeak, /quit\n")

    client = _init_g1_audio()
    lang = "ar"
    backend = args.backend
    voice = args.voice

    while True:
        try:
            text = input(f"[{lang}] > ").strip()
        except (KeyboardInterrupt, EOFError):
            print("\nBye.")
            break

        if not text:
            continue

        if text.startswith("/"):
            parts = text.split()
            cmd = parts[0]
            if cmd == "/quit":
                break
            elif cmd == "/lang" and len(parts) > 1:
                lang = parts[1]
                print(f"  Language: {lang}")
            elif cmd == "/voice" and len(parts) > 1:
                voice = parts[1]
                print(f"  Voice: {voice}")
            elif cmd == "/backend" and len(parts) > 1:
                backend = parts[1]
                print(f"  Backend: {backend}")
            elif cmd == "/save" and len(parts) > 1:
                save_path = " ".join(parts[1:])
                print(f"  Next speech will be saved to: {save_path}")
            else:
                print("  Commands: /lang ar|en, /voice <name>, /backend piper|espeak, /quit")
            continue

        try:
            audio, duration = synthesize(text, backend=backend, lang=lang, voice=voice)
            print(f"  {duration:.1f}s — playing...", end=" ", flush=True)
            _play_on_g1(client, audio)
            print("done")
        except Exception as e:
            print(f"  Error: {e}")


def cmd_list_voices(args):
    """List available Piper voices."""
    import subprocess

    print("Piper Arabic voices:")
    print("  ar_JO-kareem-low      — Jordanian Arabic (low quality, fastest)")
    print("  ar_JO-kareem-medium   — Jordanian Arabic (medium quality, recommended)")
    print()
    print("Piper English voices:")
    print("  en_US-lessac-medium   — US English male")
    print("  en_US-amy-medium      — US English female")
    print("  en_GB-alba-medium     — British English")
    print()
    print("espeak-ng Arabic voices:")
    try:
        result = subprocess.run(["espeak-ng", "--voices=ar"], capture_output=True, text=True)
        print(result.stdout)
    except FileNotFoundError:
        print("  espeak-ng not installed")


def main():
    parser = argparse.ArgumentParser(
        description="G1 Arabic TTS — Offline text-to-speech on G1 speaker",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s "مرحبا بكم في لوتاه"
  %(prog)s "مرحبا" --save marhaba.wav
  %(prog)s "Hello" --lang en
  %(prog)s --interactive
  %(prog)s --list-voices
  %(prog)s "مرحبا" --backend espeak
        """,
    )
    parser.add_argument("text", nargs="?", help="Text to speak")
    parser.add_argument("--lang", "-l", default="ar", help="Language: ar, en (default: ar)")
    parser.add_argument("--backend", "-b", default="auto", choices=["auto", "piper", "espeak"],
                        help="TTS backend (default: auto)")
    parser.add_argument("--voice", "-V", default=None, help="Piper voice name")
    parser.add_argument("--save", "-s", default=None, help="Save WAV to path")
    parser.add_argument("--no-play", action="store_true", help="Don't play on G1, just save")
    parser.add_argument("--interactive", "-i", action="store_true", help="Interactive mode")
    parser.add_argument("--list-voices", action="store_true", help="List available voices")

    args = parser.parse_args()

    if args.list_voices:
        cmd_list_voices(args)
    elif args.interactive:
        cmd_interactive(args)
    elif args.text:
        cmd_speak(args)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()