"""Orpheus TTS via Ollama + SNAC decoder.

Sends text to Ollama's Orpheus model via OpenAI-compatible completions API,
parses the <custom_token_N> responses, and decodes through SNAC to 24kHz WAV.

Voices: tara, leah, jess, leo, dan, mia, zac, zoe
Emotion tags: <laugh>, <chuckle>, <sigh>, <cough>, <gasp>, <yawn>, <groan>
"""

import re
import subprocess
import sys
import tempfile
import time
import wave

import numpy as np
import requests
import torch

SAMPLE_RATE = 24000
SNAC_CODEBOOK_SIZE = 4096

OLLAMA_URL = "http://127.0.0.1:11434"
MODEL = "legraphista/Orpheus:3b-ft-q4_k_m"

TOKEN_PATTERN = re.compile(r'<custom_token_(\d+)>')

# Cache the SNAC model across calls
_snac_model = None


def _get_snac():
    """Load and cache the SNAC decoder model."""
    global _snac_model
    if _snac_model is None:
        from snac import SNAC
        device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"  Loading SNAC decoder on {device}...", file=sys.stderr)
        _snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
    return _snac_model


def turn_token_into_id(token_str: str, index: int) -> int:
    """Convert a <custom_token_N> string to a SNAC codebook ID.

    Applies position-dependent offset: each of the 7 tokens per frame
    maps to a different SNAC codebook layer at a different offset.
    """
    match = TOKEN_PATTERN.search(token_str)
    if not match:
        return -1
    raw_id = int(match.group(1))
    return raw_id - 10 - ((index % 7) * SNAC_CODEBOOK_SIZE)


def tokens_to_audio(token_strings: list[str]) -> np.ndarray | None:
    """Convert Orpheus custom token strings to audio via SNAC.

    Redistributes the flat token stream into SNAC's 3 codebook layers:
      - codes_0: 1 per frame (positions 0)     -> coarse
      - codes_1: 2 per frame (positions 1, 4)  -> mid
      - codes_2: 4 per frame (positions 2,3,5,6) -> fine
    """
    num_frames = len(token_strings) // 7
    if num_frames == 0:
        return None

    token_strings = token_strings[:num_frames * 7]

    # Convert all tokens to IDs with position-dependent offsets
    ids = [turn_token_into_id(t, i) for i, t in enumerate(token_strings)]

    # Check for invalid tokens
    invalid = sum(1 for x in ids if x < 0 or x >= SNAC_CODEBOOK_SIZE)
    if invalid > 0:
        print(f"  Warning: {invalid}/{len(ids)} invalid token IDs", file=sys.stderr)

    # Redistribute into 3 SNAC layers
    codes_0, codes_1, codes_2 = [], [], []
    for i in range(num_frames):
        b = i * 7
        codes_0.append(ids[b + 0])
        codes_1.append(ids[b + 1])
        codes_2.append(ids[b + 2])
        codes_2.append(ids[b + 3])
        codes_1.append(ids[b + 4])
        codes_2.append(ids[b + 5])
        codes_2.append(ids[b + 6])

    # Clamp values to valid range
    def clamp(lst):
        return [max(0, min(SNAC_CODEBOOK_SIZE - 1, x)) for x in lst]

    codes_0 = clamp(codes_0)
    codes_1 = clamp(codes_1)
    codes_2 = clamp(codes_2)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    snac = _get_snac()

    codes = [
        torch.tensor(codes_0, dtype=torch.long).unsqueeze(0).to(device),
        torch.tensor(codes_1, dtype=torch.long).unsqueeze(0).to(device),
        torch.tensor(codes_2, dtype=torch.long).unsqueeze(0).to(device),
    ]

    with torch.no_grad():
        audio = snac.decode(codes)

    return audio.squeeze().cpu().numpy()


def generate_speech(text: str, voice: str = "tara") -> np.ndarray | None:
    """Generate speech from text using Orpheus via Ollama's OpenAI-compatible API."""
    prompt = f"<|audio|>{voice}: {text}<|eot_id|>"

    print(f"  Sending to Orpheus ({voice})...", file=sys.stderr)
    t0 = time.time()

    response = requests.post(
        f"{OLLAMA_URL}/v1/completions",
        json={
            "model": MODEL,
            "prompt": prompt,
            "max_tokens": 8192,
            "temperature": 0.6,
            "top_p": 0.9,
            "stream": False,
        },
    )

    result = response.json()
    resp_text = result.get("choices", [{}])[0].get("text", "")
    gen_time = time.time() - t0

    # Extract all <custom_token_N> strings
    token_strings = TOKEN_PATTERN.findall(resp_text)
    token_strings = [f"<custom_token_{t}>" for t in token_strings]

    # Skip leading special tokens (start/header tokens with small values)
    # Audio tokens have values in range ~10-28681; tokens <10 are special
    skip = 0
    for ts in token_strings:
        m = TOKEN_PATTERN.search(ts)
        if m and int(m.group(1)) < 10:
            skip += 1
        else:
            break
    if skip > 0:
        print(f"  Skipping {skip} leading special tokens", file=sys.stderr)
        token_strings = token_strings[skip:]

    print(
        f"  Got {len(token_strings)} audio tokens ({len(token_strings)//7} frames) "
        f"in {gen_time:.1f}s",
        file=sys.stderr,
    )

    if len(token_strings) < 7:
        print(f"  Insufficient tokens. Response: {resp_text[:300]}", file=sys.stderr)
        return None

    return tokens_to_audio(token_strings)


def speak(text: str, voice: str = "tara"):
    """Generate speech and play it through PipeWire."""
    audio = generate_speech(text, voice)
    if audio is None:
        print("Failed to generate audio", file=sys.stderr)
        return

    # Normalize to 16-bit PCM
    peak = max(abs(audio.max()), abs(audio.min()), 1e-8)
    audio_int16 = (audio / peak * 32767).astype(np.int16)

    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False, prefix="orpheus-")
    with wave.open(tmp.name, "wb") 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"  Audio: {duration:.1f}s at {SAMPLE_RATE}Hz -> {tmp.name}", file=sys.stderr)
    subprocess.run(["pw-play", tmp.name])


if __name__ == "__main__":
    text = " ".join(sys.argv[1:]) or (
        "Hey Ryan, this is Orpheus speaking from the fix TTS project. "
        "Pretty wild that a language model can sound this natural, right?"
    )
    speak(text)