"""Kokoro-82M TTS Space: Gradio UI + a streaming HTTP endpoint.

- Gradio UI and its API (/generate_tts) are unchanged, so existing clients
  (gradio_client.predict) keep working.
- POST /tts/stream returns raw PCM16 mono audio WHILE it is being generated:
  the text is cut into pieces (the first sentence at its first clause, so the
  first audio is short), each piece is synthesized and sent as soon as it is
  ready. No Gradio queue, no WAV file to download, resampled to 16 kHz here.

Engine (KOKORO_ENGINE): "onnx" (default) runs the same Kokoro-82M v1.0
weights through ONNX Runtime (kokoro-onnx), with phonemes from Kokoro's own
G2P (misaki) so pronunciation is unchanged; "torch" is the original
PyTorch KPipeline. If ONNX can't start, it falls back to torch.

The Space runs on 2 shared vCPUs: one synthesis at a time (a lock per piece,
so concurrent requests interleave piece by piece instead of waiting for a
whole reply).
"""
import os
import threading
import time
import urllib.request

import gradio as gr
import numpy as np
import torch
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
from kokoro import KPipeline
from scipy.signal import resample_poly

import lfm3_engine
from pieces import split_for_streaming

def _cpu_quota() -> int:
    """CPUs this container may really use. os.cpu_count() reports the host
    (16 on HF cpu-basic, which gives 2 vCPUs): that many threads would fight
    over 2 cores. Reads the cgroup quota; KOKORO_THREADS overrides."""
    env = os.getenv("KOKORO_THREADS", "").strip()
    if env.isdigit() and int(env) > 0:
        return int(env)
    try:
        quota, period = open("/sys/fs/cgroup/cpu.max").read().split()[:2]
        if quota != "max":
            return max(1, int(int(quota) / int(period)))
    except (OSError, ValueError):
        pass
    try:
        q = int(open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us").read())
        p = int(open("/sys/fs/cgroup/cpu/cpu.cfs_period_us").read())
        if q > 0:
            return max(1, q // p)
    except (OSError, ValueError):
        pass
    try:
        return max(1, len(os.sched_getaffinity(0)))
    except AttributeError:
        return max(1, os.cpu_count() or 1)


THREADS = _cpu_quota()
torch.set_num_threads(THREADS)

# Optional shared secret (Space secret TTS_API_KEY): when set, /tts/stream
# requires the X-API-Key header. The Gradio UI stays public as before.
API_KEY = os.getenv("TTS_API_KEY", "").strip()
MAX_CHARS = 600
SAMPLE_RATE = 24000   # Kokoro output

_LANGS = {"a": ("af_", "am_"), "b": ("bf_", "bm_"), "p": ("pf_", "pm_")}
_pipelines: dict = {}
_pipelines_lock = threading.Lock()
_synth_lock = threading.Lock()

# ---------------------------------------------------------------- engine

ONNX_RELEASE = "https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0"
ONNX_FILES = ("kokoro-v1.0.onnx", "voices-v1.0.bin")
MODEL_DIR = os.getenv("KOKORO_MODEL_DIR", os.path.join(os.path.dirname(os.path.abspath(__file__)), "models"))
ENGINE = os.getenv("KOKORO_ENGINE", "onnx").strip().lower()
_onnx = None  # kokoro_onnx.Kokoro when the ONNX engine is active


def _start_onnx():
    import onnxruntime as ort
    from kokoro_onnx import Kokoro

    os.makedirs(MODEL_DIR, exist_ok=True)
    for name in ONNX_FILES:
        path = os.path.join(MODEL_DIR, name)
        if not os.path.exists(path) or os.path.getsize(path) == 0:
            t0 = time.perf_counter()
            urllib.request.urlretrieve(f"{ONNX_RELEASE}/{name}", path + ".part")
            os.replace(path + ".part", path)
            print(f"[TTS] downloaded {name} in {time.perf_counter() - t0:.1f}s")
    opts = ort.SessionOptions()
    opts.intra_op_num_threads = THREADS
    opts.inter_op_num_threads = 1
    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    session = ort.InferenceSession(
        os.path.join(MODEL_DIR, ONNX_FILES[0]), sess_options=opts, providers=["CPUExecutionProvider"]
    )
    return Kokoro.from_session(session, os.path.join(MODEL_DIR, ONNX_FILES[1]))


if ENGINE == "onnx":
    try:
        _onnx = _start_onnx()
        print(f"[TTS] engine: onnx ({THREADS} threads)")
    except Exception as e:
        print(f"[TTS] onnx failed ({e}); falling back to torch")
        ENGINE = "torch"
if ENGINE != "onnx":
    print(f"[TTS] engine: torch ({THREADS} threads)")


def get_pipeline_for_voice(voice_id: str) -> KPipeline:
    """American English is loaded at start-up; British / Portuguese on first
    use. With ONNX the pipeline is G2P only (model=False)."""
    lang = next((code for code, prefixes in _LANGS.items() if voice_id.startswith(prefixes)), "a")
    with _pipelines_lock:
        if lang not in _pipelines:
            _pipelines[lang] = KPipeline(lang_code=lang, repo_id="hexgrad/Kokoro-82M", model=_onnx is None)
        return _pipelines[lang]


get_pipeline_for_voice("af_heart")


def _synth(text: str, voice_id: str, speed: float) -> np.ndarray:
    """Float32 mono audio at 24 kHz for one piece of text."""
    pipeline = get_pipeline_for_voice(voice_id)
    segments = []
    with _synth_lock:
        if _onnx is not None:
            for result in pipeline(text, voice=voice_id, split_pattern=None):
                if result.phonemes:
                    audio, _sr = _onnx.create(result.phonemes, voice=voice_id, speed=speed, is_phonemes=True)
                    segments.append(np.asarray(audio, dtype=np.float32))
        else:
            for _gs, _ps, audio in pipeline(text, voice=voice_id, speed=speed, split_pattern=None):
                seg = np.asarray(audio, dtype=np.float32)
                if seg.ndim > 1:
                    seg = seg.mean(axis=-1)
                if seg.size:
                    segments.append(seg)
    return np.concatenate(segments) if segments else np.zeros(0, dtype=np.float32)


def _to_pcm16(audio: np.ndarray, rate: int) -> bytes:
    if rate == 16000:
        audio = resample_poly(audio, 2, 3)  # 24 kHz -> 16 kHz
    return (np.clip(audio, -1.0, 1.0) * 32767).astype(np.int16).tobytes()


# ---------------------------------------------------------------- Gradio UI


def generate_tts(text: str, voice_id: str, speed: float):
    try:
        text = (text or "").strip()[:MAX_CHARS]
        if not text:
            return None
        audio = np.concatenate([_synth(p, voice_id, speed) for p in split_for_streaming(text)])
        return (SAMPLE_RATE, audio) if audio.size else None
    except Exception as e:
        print(f"[TTS] Erro: {e}")
        return None


VOICES = [
    "af_heart", "af_alloy", "af_aoede", "af_bella", "af_jessica",
    "af_kore", "af_nicole", "af_nova", "af_river", "af_sarah", "af_sky",
    "am_adam", "am_echo", "am_eric", "am_fenrir", "am_liam",
    "am_michael", "am_onyx", "am_orion", "am_puck",
    "bf_alice", "bf_emma", "bf_isabella",
    "bm_george", "bm_lewis",
    "pf_dora", "pm_alex", "pm_santa",
]

demo = gr.Interface(
    fn=generate_tts,
    inputs=[
        gr.Textbox(label="Texto", lines=3, placeholder="Digite o texto aqui..."),
        gr.Dropdown(choices=VOICES, value="af_heart", label="Voz"),
        gr.Slider(minimum=0.5, maximum=2.0, value=1.0, step=0.1, label="Velocidade"),
    ],
    outputs=gr.Audio(label="Audio Gerado"),
    title="Kokoro-82M TTS",
    description="Ingles Americano / Britanico e Portugues Brasileiro.",
)

# ---------------------------------------------------------------- HTTP API

app = FastAPI()


@app.get("/tts/health")
def tts_health():
    return {"ok": True, "engine": ENGINE, "threads": THREADS, "voices": len(VOICES), "loaded": sorted(_pipelines)}


@app.post("/tts/stream")
async def tts_stream(request: Request):
    """Body: {"text", "voice"="af_heart", "speed"=1.0, "sample_rate"=16000}.
    Response: raw PCM16 mono little-endian, streamed piece by piece."""
    if API_KEY and request.headers.get("x-api-key") != API_KEY:
        raise HTTPException(status_code=401, detail="invalid api key")
    try:
        body = await request.json()
    except Exception:
        raise HTTPException(status_code=400, detail="json body required")
    text = str(body.get("text") or "").strip()[:MAX_CHARS]
    voice = str(body.get("voice") or "af_heart")
    if voice not in VOICES:
        raise HTTPException(status_code=400, detail="unknown voice")
    try:
        speed = min(2.0, max(0.5, float(body.get("speed") or 1.0)))
        rate = int(body.get("sample_rate") or 16000)
    except (TypeError, ValueError):
        raise HTTPException(status_code=400, detail="invalid speed or sample_rate")
    if rate not in (16000, 24000):
        raise HTTPException(status_code=400, detail="sample_rate must be 16000 or 24000")
    if not text:
        return JSONResponse({"detail": "empty text"}, status_code=400)

    pieces = split_for_streaming(text)

    def generate():
        t0 = time.perf_counter()
        for i, piece in enumerate(pieces):
            pcm = _to_pcm16(_synth(piece, voice, speed), rate)
            if i == 0:
                print(f"[TTS stream] first piece in {time.perf_counter() - t0:.2f}s ({len(piece)} chars)")
            yield pcm

    return StreamingResponse(
        generate(),  # sync generator: Starlette runs it in a worker thread
        media_type="application/octet-stream",
        headers={"X-Sample-Rate": str(rate), "X-Pieces": str(len(pieces)), "Cache-Control": "no-store"},
    )


# Optional second engine (LFM3-Audio-2B) for A/B tests: routes under /lfm3/tts,
# model loaded on first use. Off unless the Space variable LFM3_ENABLED=1.
if lfm3_engine.enabled():
    def _check_key(request: Request) -> None:
        if API_KEY and request.headers.get("x-api-key") != API_KEY:
            raise HTTPException(status_code=401, detail="invalid api key")

    app.include_router(lfm3_engine.build_router(
        engine=lfm3_engine.Lfm3Engine(THREADS),
        check_key=_check_key,
        split=split_for_streaming,
        to_pcm16=_to_pcm16,
        max_chars=MAX_CHARS,
    ))
    print("[TTS] LFM3 routes enabled at /lfm3/tts (model loads on first request)")

app = gr.mount_gradio_app(app, demo, path="/")


def _warmup():
    try:
        _synth("Hello there.", "af_heart", 1.0)
        print("[TTS] warm-up done")
    except Exception as e:
        print(f"[TTS] warm-up failed: {e}")


threading.Thread(target=_warmup, daemon=True).start()

if __name__ == "__main__":
    # HF Spaces (Gradio SDK) expose GRADIO_SERVER_PORT (7860); 7861 is taken
    # by the platform there, so PORT is not used.
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("GRADIO_SERVER_PORT", "7860")))