"""Decho — public demo Space of Dibachain's Persian text-to-speech (CPU).

* Models (private repos, read with the Space secret HF_TOKEN; never exposed to visitors):
    rt     DibaAi/Decho-v2-rt   int8  - real-time distilled model (no guidance, 8 steps)
    int8   DibaAi/Decho-v2      int8  - standard
    fp32   DibaAi/Decho-v2      fp32  - studio (16 steps)
* Sign in with Hugging Face (OAuth, scope inference-api): the language-model features (smart reading, podcast script
  writer, pronunciation suggestions) run on the signed-in user's own Hugging Face account / Inference Providers credit.
  Without sign-in: samples + short single-voice tries only.
* No voice cloning in this public version (add_voice is not exposed; the speaker encoder is not even downloaded).
"""
from __future__ import annotations

import base64
import hashlib
import io
import json
import os
import re
import subprocess
import sys
import tempfile
import threading
import time
import urllib.error
import urllib.request
import uuid
from collections import OrderedDict, defaultdict, deque

import numpy as np
import soundfile as sf
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import FileResponse, JSONResponse, Response
from fastapi.staticfiles import StaticFiles
from huggingface_hub import attach_huggingface_oauth, parse_huggingface_oauth, snapshot_download
from pydantic import BaseModel

HERE = os.path.dirname(os.path.abspath(__file__))
TOKEN = os.environ.get("HF_TOKEN")                       # Space secret: read access to the private model repos
THREADS = int(os.environ.get("THREADS", "2"))
MAIN_REPO = os.environ.get("MAIN_REPO", "DibaAi/Decho-v2")
RT_REPO = os.environ.get("RT_REPO", "DibaAi/Decho-v2-rt")
MAIN_DIR = os.path.expanduser("~/model_main")
RT_DIR = os.path.expanduser("~/model_rt")
LLM_MODEL = os.environ.get("LLM_MODEL", "deepseek-ai/DeepSeek-V4.1-Flash")
ROUTER = "https://router.huggingface.co/v1/chat/completions"
BASE_FILES = ["config.json", "tokenizer.json", "decho/**", "g2p/**", "voices/**", "onnx/*.onnx", "onnx/*.npy"]

# limits: anonymous visitors get a short try, signed-in users the full studio
LIMITS = {"anon": {"chars": 220, "chars_per_hour": 2000, "streams": 1},
          "user": {"chars": 4000, "chars_per_hour": 40000, "streams": 2}}
MODELS = {"rt": {"steps": 8, "cfg": None}, "int8": {"steps": 8, "cfg": 3.0}, "fp32": {"steps": 16, "cfg": 3.0}}

state = {"ready": False, "error": None, "models": {}, "loading": {}, "started": time.time()}
_load_lock = threading.Lock()
_synth_lock = threading.Lock()
_waiting = [0]
_waiting_lock = threading.Lock()

app = FastAPI(title="Decho — Persian text-to-speech", docs_url=None, redoc_url=None)
if os.environ.get("OAUTH_CLIENT_ID"):
    attach_huggingface_oauth(app)
else:  # local run without OAuth: the app still works, signed-out
    from starlette.middleware.sessions import SessionMiddleware

    app.add_middleware(SessionMiddleware, secret_key="local-dev")


# ----------------------------------------------------------------------------------------------- models
def _download(kind: str) -> str:
    # real files, not cache symlinks: onnxruntime refuses external data that resolves outside the model folder
    if kind == "rt":
        return snapshot_download(RT_REPO, token=TOKEN, allow_patterns=BASE_FILES + ["onnx/int8/*"], local_dir=RT_DIR)
    return snapshot_download(MAIN_REPO, token=TOKEN, allow_patterns=BASE_FILES + [f"onnx/{kind}/*"], local_dir=MAIN_DIR)


def get_model(kind: str):
    if kind not in MODELS:
        raise HTTPException(400, "unknown model")
    if kind in state["models"]:
        return state["models"][kind]
    with _load_lock:
        if kind not in state["models"]:
            t0 = time.time()
            state["loading"][kind] = True
            try:
                d = _download(kind)
                if kind != "rt" and d not in sys.path:   # one runtime package per process: the main repo's (it reads the
                    sys.path.insert(0, d)                 # real-time model's config: film_layers, cfg_fixed, chunk sizes)
                from decho import Decho

                state["models"][kind] = Decho(d, precision="int8" if kind == "rt" else kind, threads=THREADS)
                state.setdefault("load_s", {})[kind] = round(time.time() - t0, 1)
            finally:
                state["loading"].pop(kind, None)
    return state["models"][kind]


def _warm():
    try:
        m = get_model("int8")                 # first: provides the runtime package
        m.synth("سلام.", steps=4, reader=False)
        r = get_model("rt")
        r.synth("سلام.", reader=False)
        state["ready"] = True
        get_model("fp32")                     # studio model in the background (large download)
    except Exception as e:  # noqa: BLE001 - logged in full, a short notice in the UI instead of a dead page
        print("model load failed:", type(e).__name__, str(e)[:500], flush=True)
        state["error"] = "model_unavailable"


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


# ----------------------------------------------------------------------------------------------- identity, limits
def who(request: Request) -> dict:
    """Signed-in Hugging Face user (OAuth) or an anonymous visitor (hashed client address)."""
    try:
        info = parse_huggingface_oauth(request)
    except Exception:  # noqa: BLE001
        info = None
    if info is not None and info.user_info is not None and info.access_token:
        u = info.user_info
        return {"kind": "user", "id": "u:" + (u.preferred_username or u.sub or "?"), "username": u.preferred_username,
                "name": u.name or u.preferred_username, "avatar": u.picture, "token": info.access_token,
                "pro": bool(u.is_pro)}
    ip = (request.headers.get("x-forwarded-for") or (request.client.host if request.client else "?")).split(",")[0].strip()
    return {"kind": "anon", "id": "a:" + hashlib.sha256(ip.encode()).hexdigest()[:16]}


_usage: dict = defaultdict(deque)   # id -> deque[(t, chars)]
_usage_lock = threading.Lock()


def charge(ident: dict, chars: int) -> None:
    lim = LIMITS[ident["kind"]]
    now = time.time()
    with _usage_lock:
        q = _usage[ident["id"]]
        while q and now - q[0][0] > 3600:
            q.popleft()
        used = sum(c for _, c in q)
        if used + chars > lim["chars_per_hour"]:
            if ident["kind"] == "anon":
                raise HTTPException(429, "limit_anon")
            raise HTTPException(429, "limit_user")
        q.append((now, chars))


_readers: "OrderedDict[str, object]" = OrderedDict()


def user_reader(ident: dict):
    """LLM reading layer on the signed-in user's own token (their Inference Providers credit)."""
    if ident["kind"] != "user":
        return None
    from decho.llm_reader import HFRouterReader

    key = ident["id"] + ":" + hashlib.sha256(ident["token"].encode()).hexdigest()[:8]
    if key not in _readers:
        _readers[key] = HFRouterReader(token=ident["token"], timeout=60.0)
        while len(_readers) > 200:
            _readers.popitem(last=False)
    _readers.move_to_end(key)
    return _readers[key]


def clean_lexicon(lex) -> dict:
    """Personal pronunciation dictionary sent by the browser: {word: Persian respelling | "r:<model text>"}."""
    out = {}
    if not isinstance(lex, dict):
        return out
    for k, v in list(lex.items())[:300]:
        k, v = str(k).strip()[:40], str(v).strip()[:80]
        if not k or not v or not re.search(r"[؀-ۿ]", k):
            continue
        if v.startswith("r:") and not re.fullmatch(r"r:[A-Za-z' ]{1,78}", v):
            continue
        out[k] = v
    return out


def wav_bytes(wav: np.ndarray, sr: int) -> bytes:
    buf = io.BytesIO()
    sf.write(buf, np.clip(np.asarray(wav, np.float32), -1, 1), sr, format="WAV", subtype="PCM_16")
    return buf.getvalue()


# ----------------------------------------------------------------------------------------------- status / account
@app.get("/api/status")
def status(request: Request):
    ident = who(request)
    m = state["models"].get("int8")
    voices = []
    if m is not None:
        meta = {}
        try:
            meta = json.load(open(os.path.join(MAIN_DIR, "voices", "voices.json"), encoding="utf-8"))
        except Exception:  # noqa: BLE001
            pass
        voices = [{"name": v, "gender": (meta.get(v) or {}).get("gender"), "fa": (meta.get(v) or {}).get("description", "")}
                  for v in m.voices]
    with _waiting_lock:
        queue = _waiting[0]
    return JSONResponse({"ready": state["ready"], "error": state["error"], "models": sorted(state["models"]),
                         "loading": sorted(state["loading"]), "voices": voices, "queue": queue,
                         "limits": LIMITS[ident["kind"]], "signed_in": ident["kind"] == "user",
                         "oauth": bool(os.environ.get("OAUTH_CLIENT_ID")),
                         "space_host": os.environ.get("SPACE_HOST", "")})


@app.get("/api/me")
def me(request: Request):
    ident = who(request)
    if ident["kind"] != "user":
        return {"signed_in": False}
    return {"signed_in": True, "username": ident["username"], "name": ident["name"], "avatar": ident["avatar"],
            "pro": ident["pro"]}


# ----------------------------------------------------------------------------------------------- synthesis (streamed)
class StreamReq(BaseModel):
    script: str | list
    model: str = "rt"                   # rt | int8 | fp32
    voice: str | None = "default"
    speed: float = 1.0
    seed: int | None = 0
    smart: bool = False                 # LLM reading layer (signed-in users, their own credit)
    register: str = "formal"            # formal | casual
    live: bool = False                  # short first chunk = earliest first audio
    pause_scale: float | None = None
    lexicon: dict | None = None


_streams: dict = {}
_streams_lock = threading.Lock()


@app.post("/api/stream/start")
def stream_start(r: StreamReq, request: Request):
    ident = who(request)
    if not state["ready"]:
        raise HTTPException(503, state["error"] or "loading")
    lim = LIMITS[ident["kind"]]
    script = r.script
    if isinstance(script, list):
        script = [{"text": str(x.get("text", ""))[:lim["chars"]], "voice": x.get("voice")} for x in script[:200]
                  if isinstance(x, dict)]
        n = sum(len(x["text"]) for x in script)
    else:
        script = str(script).strip()
        n = len(script)
    if n == 0:
        raise HTTPException(422, "empty")
    if n > lim["chars"]:
        raise HTTPException(413, "too_long_anon" if ident["kind"] == "anon" else "too_long")
    if ident["kind"] == "anon" and (isinstance(script, list) or r.model == "fp32" or r.smart):
        raise HTTPException(401, "sign_in")
    with _streams_lock:
        mine = [k for k, v in _streams.items() if v["owner"] == ident["id"] and not v["done"]]
        for k in mine[: max(0, len(mine) - lim["streams"] + 1)]:   # a new request replaces the oldest open one
            _streams.pop(k, None)
    charge(ident, n)
    m = get_model(r.model)
    kw = dict(steps=MODELS[r.model]["steps"], speed=max(0.7, min(1.4, float(r.speed))), seed=r.seed, voice=r.voice or None,
              reader=False, register="casual" if r.register == "casual" else "formal",
              lexicon=clean_lexicon(r.lexicon) or None)
    if MODELS[r.model]["cfg"] is not None:
        kw["cfg"] = MODELS[r.model]["cfg"]
    if r.live:
        kw["first_chars"] = 40
    if r.pause_scale:
        kw["pause_scale"] = max(0.5, min(2.0, float(r.pause_scale)))
    if r.smart:
        rd = user_reader(ident)
        if rd is not None:
            kw["reader"] = rd
    if kw["voice"] and kw["voice"] not in m.voices:
        raise HTTPException(422, "unknown voice")
    sid = uuid.uuid4().hex
    with _streams_lock:
        for k in [k for k, v in _streams.items() if time.time() - v["t"] > 900]:
            _streams.pop(k, None)
        _streams[sid] = {"gen": m.stream(script, **kw), "t": time.time(), "done": False, "owner": ident["id"]}
    return {"id": sid}


@app.get("/api/stream/{sid}/next")
def stream_next(sid: str, request: Request):
    st = _streams.get(sid)
    if not st or st["owner"] != who(request)["id"]:
        raise HTTPException(404, "unknown or expired stream")
    if st["done"]:
        return Response(status_code=204)
    with _waiting_lock:
        _waiting[0] += 1
    with _synth_lock:    # one synthesis at a time: the CPU is shared
        with _waiting_lock:
            _waiting[0] -= 1
        t0 = time.time()
        try:
            chunk = next(st["gen"])
        except StopIteration:
            st["done"] = True
            _streams.pop(sid, None)
            return Response(status_code=204)
        dt = time.time() - t0
    st["t"] = time.time()
    wav, sr = np.asarray(chunk["audio"], np.float32), int(chunk["sr"])
    dur = len(wav) / sr if sr else 0.0
    meta = {k: chunk.get(k) for k in ("text", "voice", "model_text", "gap_after_s", "reader", "fixed_text", "uncertain")}
    meta["model_text"] = re.sub(r"^(?:<[^>]+>\s*)+", "", meta.get("model_text") or "")
    return Response(wav_bytes(wav, sr), media_type="audio/wav", headers={
        "X-Synth-Seconds": f"{dt:.2f}", "X-Audio-Seconds": f"{dur:.2f}",
        "X-Meta": base64.b64encode(json.dumps(meta, ensure_ascii=False).encode()).decode(),
        "Access-Control-Expose-Headers": "X-Synth-Seconds, X-Audio-Seconds, X-Meta", "Cache-Control": "no-store"})


@app.post("/api/stream/{sid}/stop")
def stream_stop(sid: str, request: Request):
    st = _streams.get(sid)
    if st and st["owner"] == who(request)["id"]:
        _streams.pop(sid, None)
    return {"ok": True}


# ----------------------------------------------------------------------------------------------- language-model features
def _llm(ident: dict, messages: list, max_tokens: int, temperature: float) -> str:
    if ident["kind"] != "user":
        raise HTTPException(401, "sign_in")
    body = {"model": LLM_MODEL, "messages": messages, "max_tokens": max_tokens, "temperature": temperature,
            "chat_template_kwargs": {"thinking": False, "enable_thinking": False}}
    req = urllib.request.Request(ROUTER, data=json.dumps(body).encode(), method="POST",
                                 headers={"Authorization": f"Bearer {ident['token']}", "Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=120) as f:
            return json.loads(f.read())["choices"][0]["message"]["content"] or ""
    except urllib.error.HTTPError as e:
        if e.code in (401, 403):
            raise HTTPException(403, "hf_permission")
        if e.code == 402:
            raise HTTPException(402, "hf_credits")
        raise HTTPException(502, "llm_failed")
    except Exception:  # noqa: BLE001
        raise HTTPException(502, "llm_failed")


POD_SYS = """تو نویسندهٔ حرفه‌ای پادکست فارسی هستی. متنی می‌نویسی که برای خوانده‌شدن با صدای انسان‌گونه ساخته شده.
قالب خروجی، بدون هیچ چیز دیگر: هر خط «نام گوینده: متن». نه عنوان، نه توضیح، نه markdown، نه شماره‌گذاری.
قواعد:
- فقط از نام‌های گوینده‌ای که داده شده استفاده کن.
- نوبت‌ها کوتاه و طبیعی باشد (۱ تا ۳ جمله)، مثل گفت‌وگوی واقعی: واکنش‌های کوتاه («آره دقیقاً»، «جالبه!»، «واقعاً؟»)، سؤال، شوخی ملایم.
- گاهی خندهٔ کوتاه با «هاهاها» یا مکث فکر با «اممم…» بگذار؛ زیاده‌روی نکن.
- بین بخش‌های اصلی یک خط جدا فقط با [pause:1.0] بگذار.
- شروع: خوشامد به شنونده‌ها و معرفی نام پادکست و موضوع. پایان: جمع‌بندی کوتاه و خداحافظی گرم.
- اعداد و تاریخ را به رقم بنویس. ایموجی نگذار. جمله‌ها را با نقطه، ویرگول، علامت سؤال و تعجب درست تمام کن.
- محاوره را با املای استاندارد و رایج بنویس (مثلاً «می‌شه»، «اومدیم»، «نمی‌دونم») و نیم‌فاصله را رعایت کن. هیچ اِعرابی نگذار."""


class PodReq(BaseModel):
    source: str
    name: str = "پادکست"
    hosts: list
    tone: str = "casual"
    minutes: int = 2


@app.post("/api/podcast/script")
def podcast_script(r: PodReq, request: Request):
    ident = who(request)
    hosts = [h for h in (r.hosts or []) if isinstance(h, dict) and str(h.get("name", "")).strip()][:4]
    if not hosts:
        raise HTTPException(422, "no_hosts")
    src = (r.source or "").strip()[:8000]
    if not src:
        raise HTTPException(422, "empty")
    minutes = max(1, min(5, int(r.minutes)))
    tone = "خودمانی و محاوره‌ای (گفتار روزمرهٔ تهرانی، مؤدبانه)" if r.tone == "casual" else "رسمی، روان و حرفه‌ای"
    who_ = "، ".join(f"{h['name']} ({h.get('role') or 'میزبان'})" for h in hosts)
    user = (f"نام پادکست: {r.name or 'پادکست'}\nگوینده‌ها: {who_}\nلحن: {tone}\nطول تقریبی: {minutes} دقیقه "
            f"(حدود {minutes * 140} کلمه)\n" + ("متن منبع:" if len(src) > 300 else "موضوع:") + "\n" + src)
    text = _llm(ident, [{"role": "system", "content": POD_SYS}, {"role": "user", "content": user}],
                min(4000, 400 + minutes * 450), 0.8)
    names = {str(h["name"]).strip() for h in hosts}
    lines = []
    for ln in text.splitlines():
        ln = ln.strip().strip("*").strip()
        if ln.startswith("[pause"):
            lines.append(ln)
        elif ":" in ln or "：" in ln:
            ln = ln.replace("：", ":")
            if ln.split(":", 1)[0].strip() in names:
                lines.append(ln)
    return {"script": "\n".join(lines)}


class SuggestReq(BaseModel):
    word: str
    sentence: str | None = None


@app.post("/api/lexicon/suggest")
def lexicon_suggest(r: SuggestReq, request: Request):
    ident = who(request)
    if not state["ready"]:
        raise HTTPException(503, "loading")
    w = (r.word or "").strip()[:40]
    if not w:
        raise HTTPException(422, "empty")
    m = get_model("int8")
    rd = user_reader(ident)
    try:
        cands = m.suggest(w, sentence=(r.sentence or "")[:300] or None, reader=rd if rd is not None else False)
    except Exception:  # noqa: BLE001
        cands = []
    return {"candidates": cands, "llm": rd is not None}


class SayReq(BaseModel):
    reading: str
    word: str | None = None
    sentence: str | None = None
    voice: str | None = "default"


@app.post("/api/lexicon/say")
def lexicon_say(r: SayReq, request: Request):
    ident = who(request)
    if ident["kind"] != "user":
        raise HTTPException(401, "sign_in")
    if not state["ready"]:
        raise HTTPException(503, "loading")
    reading = (r.reading or "").strip()[:80]
    if not reading or (not re.search(r"[؀-ۿ]", reading) and not re.fullmatch(r"[A-Za-z' ]+", reading)):
        raise HTTPException(422, "bad_reading")
    charge(ident, 60)
    m = get_model("int8")
    sent = (r.sentence or "")[:300] if (r.sentence and r.word and r.word in r.sentence) else None
    with _synth_lock:
        wav, sr = m.say_word(reading, voice=r.voice if r.voice in m.voices else "default", sentence=sent, word=r.word)
    return Response(wav_bytes(wav, sr), media_type="audio/wav")


@app.post("/api/mp3")
async def to_mp3(request: Request):
    """Browser-joined WAV -> MP3 (ffmpeg) for a small, universal download."""
    if who(request)["kind"] != "user":
        raise HTTPException(401, "sign_in")
    data = await request.body()
    if not data or len(data) > 150 * 1024 * 1024:
        raise HTTPException(413, "too_large")
    with tempfile.TemporaryDirectory() as d:
        src, dst = os.path.join(d, "in.wav"), os.path.join(d, "out.mp3")
        open(src, "wb").write(data)
        p = subprocess.run(["ffmpeg", "-loglevel", "error", "-y", "-i", src, "-ac", "1", "-ar", "44100", "-b:a", "128k", dst])
        if p.returncode or not os.path.exists(dst):
            raise HTTPException(500, "mp3_failed")
        return Response(open(dst, "rb").read(), media_type="audio/mpeg",
                        headers={"Content-Disposition": "attachment; filename=decho.mp3"})


# ----------------------------------------------------------------------------------------------- pages
app.mount("/static", StaticFiles(directory=os.path.join(HERE, "static")), name="static")


@app.get("/")
def index():
    return FileResponse(os.path.join(HERE, "static", "index.html"), headers={"Cache-Control": "no-cache"})


@app.get("/robots.txt")
def robots():
    return Response("User-agent: *\nAllow: /\n", media_type="text/plain")


@app.get("/healthz")
def healthz():
    return {"ok": True, "ready": state["ready"]}