#!/usr/bin/env python3
"""OpenAI-compatible streaming TTS server for the s2-pro Egyptian fine-tune.

POST /v1/audio/speech   {model, input, voice, response_format: wav|pcm, stream}
GET  /v1/models         model listing
GET  /v1/voices         available voice names
GET  /health

Voices are (wav, txt) reference pairs in /opt/work/voices/<name>.{wav,txt}.
Streaming: chunked WAV (header + int16 PCM segments as they are generated).
"""

import argparse
import io
import struct
import time
from pathlib import Path

import numpy as np
import soundfile as sf
import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse, StreamingResponse
from loguru import logger
from pydantic import BaseModel

MODEL_ID = "s2pro-egy"
VOICES_DIR = Path("/opt/work/voices")

app = FastAPI()
ENGINE = None
SAMPLE_RATE = 44100
VOICES = {}


def wav_stream_header(sample_rate: int, channels: int = 1, bits: int = 16) -> bytes:
    # RIFF header with unknown length (0xFFFFFFFF) for streaming
    byte_rate = sample_rate * channels * bits // 8
    block_align = channels * bits // 8
    return b"".join([
        b"RIFF", struct.pack("<I", 0xFFFFFFFF), b"WAVE",
        b"fmt ", struct.pack("<IHHIIHH", 16, 1, channels, sample_rate,
                             byte_rate, block_align, bits),
        b"data", struct.pack("<I", 0xFFFFFFFF),
    ])


class SpeechRequest(BaseModel):
    model: str = MODEL_ID
    input: str
    voice: str = "masry"
    response_format: str = "wav"   # wav | pcm
    # default False: return a complete finite WAV (correct header/length,
    # like the OpenAI API). True: chunked low-latency stream for voice agents.
    stream: bool = False
    sample_rate: int = 24000       # output rate; codec native is 44100
    temperature: float = 0.8
    top_p: float = 0.8
    speed: float | None = None     # accepted for OpenAI compat; unused


def resample_audio(audio: np.ndarray, src: int, dst: int) -> np.ndarray:
    if src == dst:
        return audio
    import torch
    import torchaudio.functional as AF
    t = torch.from_numpy(np.ascontiguousarray(audio, dtype=np.float32))
    return AF.resample(t, src, dst).numpy()


SENT_SPLIT_RE = None  # compiled lazily with escaped punctuation


def split_sentences(text: str, max_len: int = 140, min_len: int = 25) -> list[str]:
    """Split text at sentence punctuation into streamable pieces."""
    global SENT_SPLIT_RE
    import re
    if SENT_SPLIT_RE is None:
        # . ! ? ؟ । plus arabic comma ، and semicolon ؛ as soft breaks
        SENT_SPLIT_RE = re.compile("([.!?؟،؛…\n]+)")
    parts = SENT_SPLIT_RE.split(text)
    # stitch punctuation back onto its sentence
    sents = []
    for i in range(0, len(parts), 2):
        s = parts[i].strip()
        p = parts[i + 1] if i + 1 < len(parts) else ""
        if s:
            sents.append((s + p).strip())
    # merge pieces smaller than min_len; cap around max_len.
    # the FIRST piece flushes early (>=20 chars, any break) for low TTFC.
    out = []
    buf = ""
    for s in sents:
        cand = (buf + " " + s).strip() if buf else s
        if not out and len(cand) >= 20:
            out.append(cand)
            buf = ""
            continue
        if len(cand) < min_len:
            buf = cand
        elif len(cand) <= max_len:
            buf = cand
            # hard sentence end -> flush
            if cand[-1] in ".!?؟…":
                out.append(buf)
                buf = ""
        else:
            if buf:
                out.append(buf)
            buf = s
    if buf:
        out.append(buf)
    return out or [text]


def load_voices():
    from fish_speech.utils.schema import ServeReferenceAudio
    VOICES.clear()
    for wav in sorted(VOICES_DIR.glob("*.wav")):
        txt = wav.with_suffix(".txt")
        if txt.exists():
            VOICES[wav.stem] = ServeReferenceAudio(
                audio=wav.read_bytes(),
                text=txt.read_text(encoding="utf-8").strip(),
            )
    logger.info(f"voices loaded: {list(VOICES)}")


@app.get("/health")
def health():
    return {"status": "ok", "model": MODEL_ID, "voices": list(VOICES)}


@app.get("/v1/models")
def models():
    return {"object": "list",
            "data": [{"id": MODEL_ID, "object": "model", "owned_by": "olimi"}]}


@app.get("/v1/voices")
def voices():
    return {"voices": list(VOICES)}


@app.post("/v1/audio/speech")
def speech(req: SpeechRequest):
    from fish_speech.utils.schema import ServeTTSRequest

    if req.voice not in VOICES:
        raise HTTPException(400, f"unknown voice '{req.voice}'; have {list(VOICES)}")
    if req.response_format not in ("wav", "pcm"):
        raise HTTPException(400, "response_format must be wav or pcm")

    treq = ServeTTSRequest(
        text=req.input,
        format="wav",
        references=[VOICES[req.voice]],
        streaming=True,
        normalize=False,          # arabic text; upstream normalizer is en/zh
        max_new_tokens=2048,
        chunk_length=100 if req.stream else 200,
        top_p=req.top_p,
        temperature=req.temperature,
        use_memory_cache="on",    # cache reference encoding between calls
    )

    t0 = time.time()

    if not req.stream:
        # complete finite file (correct header + length), OpenAI-style
        parts = []
        final_audio = None
        for res in ENGINE.inference(treq):
            if res.code == "error":
                raise HTTPException(500, f"engine error: {res.error}")
            if res.audio is None:
                continue
            sr, audio = res.audio
            if audio is None or np.size(audio) == 0:
                continue
            if res.code == "segment":
                parts.append(audio)
            elif res.code == "final":
                final_audio = audio
        audio = final_audio if final_audio is not None else (
            np.concatenate(parts) if parts else None)
        if audio is None:
            raise HTTPException(500, "no audio generated")
        logger.info(f"non-stream done {len(audio)/SAMPLE_RATE:.2f}s in {time.time()-t0:.2f}s")
        audio = resample_audio(audio, SAMPLE_RATE, req.sample_rate)
        if req.response_format == "pcm":
            body = (np.clip(audio, -1, 1) * 32767).astype("<i2").tobytes()
            media = "audio/pcm"
        else:
            buf = io.BytesIO()
            sf.write(buf, audio, req.sample_rate, format="WAV", subtype="PCM_16")
            body = buf.getvalue()
            media = "audio/wav"
        from fastapi import Response
        return Response(content=body, media_type=media)

    def synth_piece(piece_text: str):
        """Run one engine pass for a text piece, return float audio or None."""
        preq = treq.model_copy(update={"text": piece_text})
        parts, final_audio = [], None
        for res in ENGINE.inference(preq):
            if res.code == "error":
                logger.error(f"engine error: {res.error}")
                return None
            if res.audio is None:
                continue
            sr, audio = res.audio
            if audio is None or np.size(audio) == 0:
                continue
            if res.code == "segment":
                parts.append(audio)
            elif res.code == "final":
                final_audio = audio
        if final_audio is not None:
            return final_audio
        return np.concatenate(parts) if parts else None

    pieces = split_sentences(req.input)
    logger.info(f"stream: {len(pieces)} pieces")

    def gen():
        if req.response_format == "wav":
            yield wav_stream_header(req.sample_rate)
        for i, piece in enumerate(pieces):
            audio = synth_piece(piece)
            if audio is None:
                continue
            audio = resample_audio(audio, SAMPLE_RATE, req.sample_rate)
            pcm = (np.clip(audio, -1, 1) * 32767).astype("<i2").tobytes()
            logger.info(f"piece {i+1}/{len(pieces)} {len(pcm)}B t+{time.time()-t0:.2f}s")
            yield pcm

    media = "audio/wav" if req.response_format == "wav" else "audio/pcm"
    return StreamingResponse(gen(), media_type=media)


def main():
    global ENGINE, SAMPLE_RATE
    ap = argparse.ArgumentParser()
    ap.add_argument("--model-dir", required=True)
    ap.add_argument("--port", type=int, default=8000)
    ap.add_argument("--compile", action="store_true", default=True)
    args = ap.parse_args()

    from tools.server.model_manager import ModelManager
    mm = ModelManager(
        mode="tts", device="cuda", half=False, compile=args.compile,
        llama_checkpoint_path=args.model_dir,
        decoder_checkpoint_path=str(Path(args.model_dir) / "codec.pth"),
        decoder_config_name="modded_dac_vq",
    )
    ENGINE = mm.tts_inference_engine
    SAMPLE_RATE = ENGINE.decoder_model.sample_rate
    load_voices()
    logger.info(f"ready on :{args.port} sr={SAMPLE_RATE}")
    uvicorn.run(app, host="0.0.0.0", port=args.port, log_level="warning")


if __name__ == "__main__":
    main()