"""
Test Orpheus TTS via Ollama + SNAC decoder.

Prerequisites:
  1. bash setup-tts.sh                    (one-shot: installs everything)
  -- OR manually --
  1. ollama pull sematre/orpheus:en
  2. bash download-tts-models.sh snac     (downloads SNAC via hf-mirror.com)
  3. ollama serve                          (must be running)

Usage:
  .venv-qwen-tts/bin/python test_orpheus_tts.py
"""
import os
import re
import sys
import time
import json
import struct
import wave
import urllib.request

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
SNAC_MODEL_DIR = os.path.join(SCRIPT_DIR, "models", "snac_24khz")
OLLAMA_URL = "http://localhost:11434"
MODEL = "sematre/orpheus:en"

AUDIO_TOKEN_RE = re.compile(r"<custom_token_(\d+)>")


def check_ollama():
    """Verify Ollama is running and model is available."""
    try:
        req = urllib.request.Request(f"{OLLAMA_URL}/api/tags")
        with urllib.request.urlopen(req, timeout=3) as resp:
            data = json.loads(resp.read())
            names = [m["name"] for m in data.get("models", [])]
            if not any(MODEL in n for n in names):
                print(f"ERROR: Model '{MODEL}' not found. Run: ollama pull {MODEL}")
                return False
            return True
    except Exception as e:
        print(f"ERROR: Cannot connect to Ollama at {OLLAMA_URL}: {e}")
        print("Run: ollama serve")
        return False


def check_snac():
    """Verify SNAC model is downloaded."""
    if not os.path.isdir(SNAC_MODEL_DIR):
        print(f"ERROR: SNAC decoder not found at {SNAC_MODEL_DIR}")
        print("Run: bash setup-tts.sh   (or: bash download-tts-models.sh snac)")
        return False
    return True


def load_snac():
    """Load SNAC audio codec."""
    import torch
    import snac

    print(f"Loading SNAC decoder from {SNAC_MODEL_DIR}...")
    model = snac.SNAC.from_pretrained(SNAC_MODEL_DIR)
    model.eval()
    return model


def generate_tokens(text: str, voice: str = "tara") -> str:
    """Call Ollama to generate audio tokens from text."""
    prompt = f"<custom_token_3><|begin_of_text|>{voice}: {text}<|eot_id|><custom_token_4><custom_token_5><custom_token_1>"

    payload = json.dumps({
        "model": MODEL,
        "prompt": prompt,
        "stream": False,
        "options": {
            "temperature": 0.6,
            "top_p": 0.9,
            "repeat_penalty": 1.1,
            "num_predict": 10240,
            "stop": ["<|end_of_text|>"],
        },
    }).encode()

    req = urllib.request.Request(
        f"{OLLAMA_URL}/api/generate",
        data=payload,
        headers={"Content-Type": "application/json"},
    )

    print("Generating audio tokens via Ollama...")
    t0 = time.time()
    with urllib.request.urlopen(req, timeout=120) as resp:
        result = json.loads(resp.read())

    elapsed = time.time() - t0
    response_text = result.get("response", "")
    token_count = len(AUDIO_TOKEN_RE.findall(response_text))
    print(f"Generated {token_count} audio tokens in {elapsed:.1f}s")
    return response_text


def decode_tokens(response_text: str, snac_model) -> tuple:
    """Convert audio tokens to WAV audio."""
    import torch

    tokens = AUDIO_TOKEN_RE.findall(response_text)
    if not tokens:
        print("ERROR: No audio tokens found in response")
        return None, 0

    audio_ids = [
        int(tok) - 10 - ((idx % 7) * 4096)
        for idx, tok in enumerate(tokens)
    ]

    # Trim to multiple of 7
    audio_ids = audio_ids[: len(audio_ids) // 7 * 7]
    if len(audio_ids) == 0:
        print("ERROR: Not enough audio tokens to decode")
        return None, 0

    audio_tensor = torch.tensor(audio_ids, dtype=torch.int32).reshape(-1, 7)
    codes_0 = audio_tensor[:, 0].unsqueeze(0)
    codes_1 = torch.stack((audio_tensor[:, 1], audio_tensor[:, 4])).t().flatten().unsqueeze(0)
    codes_2 = (
        torch.stack((audio_tensor[:, 2], audio_tensor[:, 3], audio_tensor[:, 5], audio_tensor[:, 6]))
        .t()
        .flatten()
        .unsqueeze(0)
    )

    print("Decoding audio...")
    with torch.inference_mode():
        audio_hat = snac_model.decode([codes_0, codes_1, codes_2])

    audio_np = audio_hat[0].squeeze().numpy()
    return audio_np, 24000


def save_wav(audio_np, sample_rate: int, path: str):
    """Save numpy audio array as 16-bit WAV."""
    import numpy as np

    # Normalize to int16
    audio_int16 = (audio_np * 32767).clip(-32768, 32767).astype(np.int16)

    with wave.open(path, "w") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sample_rate)
        wf.writeframes(audio_int16.tobytes())

    duration = len(audio_int16) / sample_rate
    print(f"Saved {path} ({duration:.1f}s, {sample_rate} Hz)")


def main():
    print("=== Orpheus TTS Test (Ollama + SNAC) ===\n")

    if not check_ollama():
        return
    if not check_snac():
        return

    snac_model = load_snac()

    # Voices: tara, leah, jess, leo, dan, mia, zac, zoe
    tests = [
        ("Hello! This is Orpheus text to speech, running entirely locally through Ollama.", "tara"),
        ("<laugh> That's amazing! Local AI speech generation without any cloud services!", "leo"),
    ]

    for i, (text, voice) in enumerate(tests):
        print(f"\n--- Test {i+1}: voice={voice} ---")
        print(f"Text: {text[:80]}...")

        response = generate_tokens(text, voice)
        audio, sr = decode_tokens(response, snac_model)

        if audio is not None:
            outpath = os.path.join(SCRIPT_DIR, f"test_orpheus_{voice}.wav")
            save_wav(audio, sr, outpath)

    print("\n=== Done! Open the .wav files to listen. ===")
    if sys.platform == "darwin":
        print("Play with:  afplay test_orpheus_tara.wav")
    else:
        print("Play with:  aplay test_orpheus_tara.wav  (or: ffplay -nodisp -autoexit test_orpheus_tara.wav)")


if __name__ == "__main__":
    main()
