"""Orpheus TTS via llama-server + streaming SNAC decoder.

Sends text to a llama-server completions endpoint with streaming enabled,
parses <custom_token_N> responses as they arrive via SSE, and decodes
through SNAC in batches of 28 tokens (4 frames) for overlapped inference.
SNAC runs on CPU; the LLM runs on GPU via llama-server.
"""

import asyncio
import json
import os
import re
import sys
import time

import httpx
import numpy as np

from ..audio import wav_duration, write_wav
from ..settings import settings
from .base import TTSEngine, TTSResult

SAMPLE_RATE = 24000
SNAC_CODEBOOK_SIZE = 4096
TOKEN_PATTERN = re.compile(r"<custom_token_(\d+)>")
TOKENS_PER_FRAME = 7

ALL_VOICES = ["tara", "leah", "jess", "leo", "dan", "mia", "zac", "zoe"]


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], snac_model) -> 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
    """
    import torch

    num_frames = len(token_strings) // 7
    if num_frames == 0:
        return None

    token_strings = token_strings[: num_frames * 7]
    ids = [_turn_token_into_id(t, i) for i, t in enumerate(token_strings)]

    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)

    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])

    def clamp(lst):
        return [max(0, min(SNAC_CODEBOOK_SIZE - 1, x)) for x in lst]

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

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

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


def _load_snac():
    """Load SNAC model on CPU. Called once, lazily.

    Sets CUDA_VISIBLE_DEVICES="" for the entire process to prevent PyTorch
    from allocating GPU memory for SNAC. Safe because all GPU inference is
    handled by llama-server in a separate container.
    """
    os.environ["CUDA_VISIBLE_DEVICES"] = ""
    from snac import SNAC

    return SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to("cpu")


class OrpheusEngine(TTSEngine):
    """Orpheus TTS via llama-server completions API + streaming SNAC decode.

    Streams tokens from llama-server via SSE, decodes in batches of 28
    tokens (4 SNAC frames) to overlap GPU inference with CPU audio decode.
    SNAC is lazy-loaded on first synthesize() call.
    """

    name = "orpheus"
    default_voice = "tara"

    def __init__(self, orpheus_url: str) -> None:
        self._snac = None
        self._snac_lock = asyncio.Lock()
        self._url = orpheus_url
        self._client = httpx.AsyncClient(
            timeout=httpx.Timeout(connect=10.0, read=600.0, write=30.0, pool=30.0),
            limits=httpx.Limits(max_connections=4, max_keepalive_connections=2),
        )

    async def close(self) -> None:
        """Clean up httpx client."""
        await self._client.aclose()

    async def _get_snac(self):
        """Lazy-load SNAC on first use."""
        if self._snac is not None:
            return self._snac
        async with self._snac_lock:
            if self._snac is not None:
                return self._snac
            print("Loading SNAC decoder on CPU (first Orpheus call)...", file=sys.stderr)
            t0 = time.time()
            loop = asyncio.get_running_loop()
            self._snac = await loop.run_in_executor(None, _load_snac)
            print(f"  SNAC ready in {time.time() - t0:.1f}s", file=sys.stderr)
            return self._snac

    async def synthesize(self, text: str, voice: str | None = None) -> TTSResult:
        voice = voice or self.default_voice
        prompt = f"<|audio|>{voice}: {text}<|eot_id|>"

        loop = asyncio.get_running_loop()
        snac = await self._get_snac()

        # Stream tokens from llama-server via SSE, then decode in one pass.
        # SNAC's convolutional decoder has a receptive field spanning multiple
        # frames, so batched decode creates boundary artifacts. Since pw-play
        # waits for the full WAV anyway, single-pass decode costs only ~2s
        # extra on a 50s generation — and produces clean audio.
        t0 = time.time()
        token_strings: list[str] = []
        total_tokens = 0
        dropped_lines = 0

        try:
            async with self._client.stream(
                "POST",
                f"{self._url}/v1/completions",
                json={
                    "prompt": prompt,
                    "max_tokens": 8192,
                    "temperature": 0.6,
                    "top_p": 0.9,
                    "stream": True,
                },
            ) as response:
                try:
                    response.raise_for_status()
                except httpx.HTTPStatusError as e:
                    body = await response.aread()
                    raise RuntimeError(
                        f"llama-server returned {e.response.status_code}: "
                        f"{body[:500].decode(errors='replace')}"
                    ) from e

                async for line in response.aiter_lines():
                    # SSE format: "data: {...}" or "data: [DONE]"
                    if not line.startswith("data: "):
                        continue
                    payload = line[6:]
                    if payload == "[DONE]":
                        break

                    try:
                        chunk = json.loads(payload)
                    except json.JSONDecodeError:
                        dropped_lines += 1
                        continue

                    chunk_text = chunk.get("choices", [{}])[0].get("text", "")
                    matches = TOKEN_PATTERN.findall(chunk_text)
                    for m in matches:
                        token_val = int(m)
                        # Skip leading special tokens (value < 10)
                        if total_tokens == 0 and token_val < 10:
                            continue
                        token_strings.append(f"<custom_token_{m}>")
                        total_tokens += 1

        except httpx.ConnectError as e:
            raise RuntimeError(f"Cannot reach llama-server at {self._url}: {e}") from e
        except (httpx.ReadError, httpx.RemoteProtocolError) as e:
            if total_tokens < TOKENS_PER_FRAME:
                raise RuntimeError(
                    f"llama-server connection lost with no usable tokens: {e}"
                ) from e
            print(
                f"  Warning: llama-server connection lost after {total_tokens} tokens. "
                "Decoding what we have.",
                file=sys.stderr,
            )

        gen_time = time.time() - t0

        if dropped_lines > 0:
            print(
                f"  Warning: {dropped_lines} SSE lines had malformed JSON",
                file=sys.stderr,
            )

        # Drop partial frame at end (at most 6 tokens / 0.29ms lost)
        usable = len(token_strings) - (len(token_strings) % TOKENS_PER_FRAME)

        num_frames = usable // TOKENS_PER_FRAME
        tok_per_sec = total_tokens / gen_time if gen_time > 0 else 0
        print(
            f"  Orpheus: {total_tokens} tokens ({num_frames} frames) "
            f"in {gen_time:.1f}s ({tok_per_sec:.1f} tok/s)",
            file=sys.stderr,
        )

        if usable < TOKENS_PER_FRAME:
            raise RuntimeError(
                f"Orpheus returned insufficient tokens ({total_tokens}). "
                "Check llama-server logs."
            )

        # Single-pass SNAC decode — full context across all frames
        audio = await loop.run_in_executor(
            None, _tokens_to_audio, token_strings[:usable], snac
        )
        if audio is None:
            raise RuntimeError("SNAC decoding produced no audio")
        path = write_wav(audio, SAMPLE_RATE, prefix="orpheus-")

        return TTSResult(
            audio_path=path,
            sample_rate=SAMPLE_RATE,
            duration_seconds=wav_duration(path),
            engine=self.name,
            voice=voice,
        )

    async def list_voices(self) -> list[str]:
        blacklist = settings.blacklisted_voices
        return sorted(v for v in ALL_VOICES if v.lower() not in blacklist)

    async def check_health(self) -> dict:
        try:
            resp = await self._client.get(f"{self._url}/health")
            resp.raise_for_status()
            data = resp.json()
            status = data.get("status", "unknown")
            return {
                "status": "healthy" if status == "ok" else "degraded",
                "engine": self.name,
                "backend": "llama-server",
                "server_status": status,
            }
        except Exception as e:
            return {"status": "unhealthy", "engine": self.name, "error": str(e)}
