import os
import re
import secrets
import threading
import uuid

from fastapi import Depends, FastAPI, HTTPException, Header, Request
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field

# strip(): pasted secrets often carry a trailing newline.
API_KEY = os.environ["API_KEY"].strip()
OUT_DIR = os.environ.get("OUT_DIR", "/tmp/audio")
os.makedirs(OUT_DIR, exist_ok=True)

app = FastAPI(title="Chatterbox TTS")
_model = None
# ponytail: global lock. chatterbox keeps per-generation state on the shared model
# (alignment_stream_analyzer), so concurrent generate() calls corrupt each other.
# Sync endpoints run in a threadpool, so "one worker" is not one caller.
# Upgrade path: a real request queue / model pool if throughput matters.
_gen_lock = threading.Lock()


def auth(x_api_key: str = Header(None)):
    if not x_api_key or not secrets.compare_digest(x_api_key.strip(), API_KEY):
        raise HTTPException(401, "bad api key")


def model():
    # ponytail: lazy singleton so the Space boots fast and healthchecks pass.
    # Worst case two requests both build it and one wins; generate() is serialized
    # by _gen_lock anyway.
    global _model
    if _model is None:
        import functools

        import torch

        # chatterbox's from_local() loads s3gen.pt without map_location, so the
        # CUDA-saved tensors blow up on CPU. Default every torch.load to CPU.
        torch.load = functools.partial(torch.load, map_location="cpu")

        from chatterbox.mtl_tts import ChatterboxMultilingualTTS

        _model = ChatterboxMultilingualTTS.from_pretrained(device="cpu")
    return _model


class TTSRequest(BaseModel):
    # NOTE: upstream generate() hardcodes max_new_tokens=1000 at S3_TOKEN_RATE=25, so
    # audio past ~40s (roughly 600-700 chars) is silently truncated, not errored.
    text: str = Field(min_length=1, max_length=1000)
    lang: str = Field("en", min_length=2, max_length=5)
    exaggeration: float = Field(0.5, ge=0.0, le=2.0)
    cfg_weight: float = Field(0.5, ge=0.0, le=1.0)


@app.post("/tts", dependencies=[Depends(auth)])
def tts(req: TTSRequest, request: Request):
    import torchaudio  # ponytail: lazy so the module imports without torch installed

    m = model()
    lang = req.lang.lower()
    if lang not in m.get_supported_languages():
        raise HTTPException(400, f"unsupported lang: {lang}")

    with _gen_lock:
        wav = m.generate(
            req.text,
            language_id=lang,
            exaggeration=req.exaggeration,
            cfg_weight=req.cfg_weight,
        )
    stem = uuid.uuid4().hex
    torchaudio.save(os.path.join(OUT_DIR, stem + ".wav"), wav, m.sr)
    with open(os.path.join(OUT_DIR, stem + ".srt"), "w") as f:
        f.write(srt(req.text, wav.shape[-1] / m.sr))
    base = str(request.base_url)
    return {
        "url": base + f"audio/{stem}.wav",
        "srt_url": base + f"audio/{stem}.srt",
        "lang": lang,
    }


def _ts(t):
    ms = round(t * 1000)
    return f"{ms // 3600000:02d}:{ms // 60000 % 60:02d}:{ms // 1000 % 60:02d},{ms % 1000:03d}"


def srt(text, duration):
    # ponytail: chatterbox gives no word timings, so split on sentence enders and
    # give each cue a slice of the duration proportional to its length. Swap in
    # forced alignment (whisperx/aeneas) if the drift ever matters.
    parts = [s.strip() for s in re.split(r"(?<=[.!?。！？])\s+", text) if s.strip()]
    total = sum(len(p) for p in parts) or 1
    out, t = [], 0.0
    for i, p in enumerate(parts, 1):
        end = t + duration * len(p) / total
        out.append(f"{i}\n{_ts(t)} --> {_ts(end)}\n{p}\n")
        t = end
    return "\n".join(out)


@app.get("/audio/{name}")
def download(name: str):
    # ponytail: uuid names only; reject anything else instead of path-joining it.
    if len(name) != 36 or name[32:] not in (".wav", ".srt") or not name[:32].isalnum():
        raise HTTPException(400, "bad name")
    path = os.path.join(OUT_DIR, name)
    if not os.path.isfile(path):
        raise HTTPException(404, "unknown file")
    media = "audio/wav" if name.endswith(".wav") else "application/x-subrip"
    return FileResponse(path, media_type=media, filename=name)


@app.get("/")
def health():
    return {"ok": True, "loaded": _model is not None}