"""OpenAI-kompatible Sprachausgabe für NEX: POST /v1/audio/speech mit Piper auf der CPU.

Schlanker Wrapper um piper-tts (Python-API) und ffmpeg:
- Formate: mp3 (Standard), opus (Ogg-Container), aac, flac, pcm (int16 LE mono) und wav.
- mp3/opus/aac/flac/pcm werden satzweise gestreamt (chunked), wav kommt vollständig mit
  korrektem Header (kompatibel mit `knaster bench tts` und dem Canary-Prober).
- pcm wird auf PCM_SAMPLE_RATE umgerechnet (Standard 24000 Hz wie bei OpenAI); die Rate steht in
  `X-Audio-Sample-Rate` und in `Content-Type: audio/pcm; rate=<Hz>`.
- Datenschutz: Text und Audio werden weder gespeichert noch geloggt; alles läuft über Pipes im RAM.

Umgebung:
  PIPER_VOICES_DIR     Verzeichnis mit <stimme>.onnx und <stimme>.onnx.json (Standard /voices)
  PIPER_VOICES         Kommagetrennte Stimmen-IDs (Standard de_DE-thorsten-high,de_DE-kerstin-low)
  PIPER_DEFAULT_VOICE  Standardstimme (Standard: erste aus PIPER_VOICES)
  PCM_SAMPLE_RATE      Rate für response_format=pcm, 0 = native Rate der Stimme (Standard 24000)
  TTS_MAX_CHARS        Höchstlänge von input in Zeichen (Standard 5000)
  TTS_PARALLEL         Gleichzeitige Synthesen (Standard 2)
"""

from __future__ import annotations

import asyncio
import io
import logging
import os
import threading
import wave
from collections.abc import AsyncGenerator, AsyncIterator
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, Response, StreamingResponse
from piper import PiperVoice, SynthesisConfig

log = logging.getLogger("nex-tts")
logging.basicConfig(level=os.environ.get("LOG_LEVEL", "WARNING").upper(),
                    format="%(asctime)s %(levelname)s %(name)s: %(message)s")

VOICES_DIR = Path(os.environ.get("PIPER_VOICES_DIR", "/voices"))
VOICE_IDS = [v.strip() for v in os.environ.get("PIPER_VOICES", "de_DE-thorsten-high,de_DE-kerstin-low").split(",")
             if v.strip()]
DEFAULT_VOICE = os.environ.get("PIPER_DEFAULT_VOICE") or VOICE_IDS[0]
PCM_RATE = int(os.environ.get("PCM_SAMPLE_RATE", "24000"))
MAX_CHARS = int(os.environ.get("TTS_MAX_CHARS", "5000"))
PARALLEL = threading.Semaphore(int(os.environ.get("TTS_PARALLEL", "2")))  # CPU nicht überbuchen

# ffmpeg-Ausgabeparameter je Format: (Argumente, Content-Type)
ENCODERS: dict[str, tuple[list[str], str]] = {
    "mp3": (["-c:a", "libmp3lame", "-b:a", "64k", "-id3v2_version", "0", "-write_xing", "0", "-f", "mp3"], "audio/mpeg"),
    # Ogg-Seiten klein halten (20 ms), sonst puffert der Muxer bis zu 1 s
    "opus": (["-ar", "48000", "-c:a", "libopus", "-b:a", "32k", "-f", "ogg", "-page_duration", "20000"], "audio/ogg"),
    "aac": (["-c:a", "aac", "-b:a", "64k", "-f", "adts"], "audio/aac"),
    "flac": (["-c:a", "flac", "-f", "flac"], "audio/flac"),
}
FORMATS = {*ENCODERS, "pcm", "wav"}

voices: dict[str, PiperVoice] = {}
aliases: dict[str, str] = {}


def load_voices() -> None:
    for vid in VOICE_IDS:
        model = VOICES_DIR / f"{vid}.onnx"
        voice = PiperVoice.load(str(model), config_path=str(model) + ".json", use_cuda=False)
        voices[vid] = voice
        # Kurzname: de_DE-thorsten-high → thorsten-high und thorsten
        short = vid.split("-", 1)[1] if "-" in vid else vid
        aliases.setdefault(short, vid)
        aliases.setdefault(short.rsplit("-", 1)[0], vid)
        # Aufwärmen, damit der erste echte Aufruf nicht die ONNX-Initialisierung bezahlt
        for _ in voice.synthesize("Hallo."):
            pass
    if DEFAULT_VOICE not in voices:
        raise RuntimeError(f"Standardstimme {DEFAULT_VOICE} ist nicht geladen")


def resolve_voice(name: str | None) -> str:
    if not name or name == "default":
        return DEFAULT_VOICE
    if name in voices:
        return name
    if name in aliases:
        return aliases[name]
    # OpenAI-Clients schicken oft "alloy" o. Ä.: nicht scheitern, sondern Standardstimme nehmen
    log.warning("Stimme %r unbekannt, nutze %s", name[:40], DEFAULT_VOICE)
    return DEFAULT_VOICE


async def synth_pcm(voice: PiperVoice, text: str, cfg: SynthesisConfig) -> AsyncIterator[bytes]:
    """Piper satzweise in einem Thread rechnen und die int16-Blöcke asynchron weiterreichen."""
    loop = asyncio.get_running_loop()
    queue: asyncio.Queue[bytes | BaseException | None] = asyncio.Queue()
    stop = threading.Event()

    def work() -> None:
        try:
            with PARALLEL:
                for chunk in voice.synthesize(text, syn_config=cfg):
                    if stop.is_set():
                        break
                    loop.call_soon_threadsafe(queue.put_nowait, chunk.audio_int16_bytes)
        except BaseException as exc:  # noqa: BLE001 – an den Aufrufer weitergeben
            loop.call_soon_threadsafe(queue.put_nowait, exc)
        finally:
            loop.call_soon_threadsafe(queue.put_nowait, None)

    loop.run_in_executor(None, work)
    try:
        while True:
            item = await queue.get()
            if item is None:
                return
            if isinstance(item, BaseException):
                raise item
            yield item
    finally:
        stop.set()  # Client weg oder fertig: Thread bricht nach dem laufenden Satz ab


async def encode(pcm: AsyncIterator[bytes], rate: int, out_args: list[str]) -> AsyncIterator[bytes]:
    """Rohes PCM per ffmpeg (stdin → stdout) umwandeln und dabei weiterstreamen."""
    proc = await asyncio.create_subprocess_exec(
        "ffmpeg", "-hide_banner", "-nostdin", "-loglevel", "error",
        "-f", "s16le", "-ar", str(rate), "-ac", "1", "-i", "pipe:0",
        *out_args, "-flush_packets", "1", "pipe:1",
        stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.DEVNULL)
    assert proc.stdin is not None and proc.stdout is not None

    async def feed() -> None:
        try:
            async for block in pcm:
                proc.stdin.write(block)
                await proc.stdin.drain()
        finally:
            proc.stdin.close()

    feeder = asyncio.create_task(feed())
    try:
        while data := await proc.stdout.read(8192):
            yield data
        await feeder
        if await proc.wait() != 0:
            raise RuntimeError(f"ffmpeg beendet mit {proc.returncode}")
    finally:
        if not feeder.done():
            feeder.cancel()
        if proc.returncode is None:
            proc.kill()
            await proc.wait()


async def primed(stream: AsyncGenerator[bytes, None]) -> AsyncIterator[bytes]:
    """Erst den ersten Audioblock abwarten, dann die Header senden.

    So misst die Zeit bis zum ersten Byte echtes Audio, und ein Fehler beim Start wird noch
    als HTTP-Fehler gemeldet. Fehler mitten im Stream werden nur ohne Text protokolliert.
    """
    try:
        first = await anext(stream)
    except StopAsyncIteration:
        first = b""
    except Exception as exc:  # noqa: BLE001
        log.error("Synthese fehlgeschlagen: %s", type(exc).__name__)
        raise HTTPException(500, "Synthese fehlgeschlagen") from None

    async def rest() -> AsyncIterator[bytes]:
        try:
            if first:
                yield first
            async for data in stream:
                yield data
        except Exception as exc:
            log.error("Synthese abgebrochen: %s", type(exc).__name__)
            raise
        finally:
            await stream.aclose()

    return rest()


@asynccontextmanager
async def lifespan(_: FastAPI) -> AsyncIterator[None]:
    load_voices()
    log.warning("nex-tts bereit: Stimmen %s, Standard %s, PCM %s Hz", ", ".join(voices), DEFAULT_VOICE,
                PCM_RATE or "nativ")
    yield


app = FastAPI(title="nex-tts (Piper)", docs_url=None, redoc_url=None, openapi_url=None, lifespan=lifespan)


@app.get("/health")
def health() -> dict:
    return {"status": "ok", "voices": list(voices), "default": DEFAULT_VOICE}


@app.get("/v1/models")
def models() -> dict:
    return {"object": "list", "data": [{"id": "piper", "object": "model", "owned_by": "knaster"}]}


@app.get("/v1/audio/voices")
def list_voices() -> dict:
    return {"voices": [{"id": vid, "sample_rate": v.config.sample_rate} for vid, v in voices.items()],
            "default": DEFAULT_VOICE}


@app.post("/v1/audio/speech")
async def speech(request: Request) -> Response:
    try:
        body = await request.json()
    except ValueError:
        raise HTTPException(400, "Body muss JSON sein") from None
    if not isinstance(body, dict):
        raise HTTPException(400, "Body muss ein JSON-Objekt sein")
    # Unbekannte Felder (model, instructions, stream_format, …) werden ignoriert.
    text = body.get("input")
    if not isinstance(text, str) or not text.strip():
        raise HTTPException(400, "Feld input fehlt oder ist leer")
    if len(text) > MAX_CHARS:
        raise HTTPException(400, f"input ist länger als {MAX_CHARS} Zeichen")
    fmt = str(body.get("response_format") or "mp3").lower()
    if fmt not in FORMATS:
        raise HTTPException(400, f"response_format {fmt[:20]!r} nicht unterstützt ({', '.join(sorted(FORMATS))})")
    try:
        speed = float(body.get("speed") or 1.0)
    except (TypeError, ValueError):
        raise HTTPException(400, "speed muss eine Zahl sein") from None
    speed = min(max(speed, 0.25), 4.0)

    vid = resolve_voice(body.get("voice") if isinstance(body.get("voice"), str) else None)
    voice = voices[vid]
    rate = voice.config.sample_rate
    cfg = SynthesisConfig(length_scale=(voice.config.length_scale or 1.0) / speed)
    headers = {"X-Voice": vid, "Cache-Control": "no-store"}

    if fmt == "wav":
        # Vollständig mit korrektem Header (kurze Texte; für Streaming pcm/mp3/opus nehmen)
        buf = io.BytesIO()
        with wave.open(buf, "wb") as w:
            w.setnchannels(1)
            w.setsampwidth(2)
            w.setframerate(rate)
            async for block in synth_pcm(voice, text, cfg):
                w.writeframes(block)
        return Response(buf.getvalue(), media_type="audio/wav", headers=headers)

    if fmt == "pcm":
        out_rate = PCM_RATE or rate
        headers |= {"X-Audio-Sample-Rate": str(out_rate), "X-Audio-Channels": "1", "X-Audio-Format": "s16le"}
        media = f"audio/pcm; rate={out_rate}"
        if out_rate == rate:
            stream = synth_pcm(voice, text, cfg)
        else:
            stream = encode(synth_pcm(voice, text, cfg), rate, ["-ar", str(out_rate), "-f", "s16le"])
        return StreamingResponse(await primed(stream), media_type=media, headers=headers)

    out_args, media = ENCODERS[fmt]
    stream = encode(synth_pcm(voice, text, cfg), rate, out_args)
    return StreamingResponse(await primed(stream), media_type=media, headers=headers)


@app.exception_handler(HTTPException)
async def _http_error(_: Request, exc: HTTPException) -> JSONResponse:
    # OpenAI-Fehlerformat, ohne den Eingabetext zu wiederholen
    return JSONResponse({"error": {"message": exc.detail, "type": "invalid_request_error"}},
                        status_code=exc.status_code)