# -*- coding: utf-8 -*-
"""وسيط «المجوّد/بيان» (HF Space) — ثلاث مهامّ:
  1) بروكسي آمن لخادمَي HF المحميّين (/stt نصّيّ، /analyze فونيميّ) — التوكن سرّ Space.
  2) حماية إساءة الاستخدام: حدّ معدّل لكلّ IP (نافذة منزلقة) + سقف جلسات بثّ متزامنة.
  3) /stream: بثّ لحظيّ (WebSocket) بنموذج Muno459 streaming (cache-aware CTC) يعمل هنا
     على CPU — إطارات int16@16k → نصّ تراكميّ + نافذة recent (~4.5ث) لكلّ قطعة 112-mel.
الثوابت الصوتيّة (نافذة/mel-fb/CMVN) من mel_const.json — نفس ملفّ العميل، تطابقٌ بتّيّ.
"""
import os, io, json, time, base64, asyncio, threading
from collections import deque, defaultdict
import numpy as np
import requests
from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse

EP = os.environ.get("MUAALEM_ENDPOINT_URL", "").rstrip("/")
TOK = os.environ.get("MUAALEM_TOKEN", "")
STT_EP = os.environ.get("STT_ENDPOINT_URL", "").rstrip("/")
STT_TOK = os.environ.get("STT_TOKEN", "") or TOK

app = FastAPI()
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

# ───────────────────────── حماية المعدّل (لكلّ IP) ─────────────────────────
_BUCKETS: dict = defaultdict(deque)
_WS_LIVE: dict = defaultdict(int)
_LOCK = threading.Lock()

def _client_ip(scope_headers, client):
    for k, v in scope_headers:
        if k == b"x-forwarded-for":
            return v.decode().split(",")[0].strip()
    return client.host if client else "?"

def _allow(ip: str, route: str, limit: int, per: int = 60) -> bool:
    with _LOCK:
        q = _BUCKETS[(ip, route)]
        now = time.time()
        while q and now - q[0] > per:
            q.popleft()
        if len(q) >= limit:
            return False
        q.append(now)
        return True

def _ip_http(req: Request) -> str:
    return (req.headers.get("x-forwarded-for", "").split(",")[0].strip()
            or (req.client.host if req.client else "?"))

# ───────────────────────── محرّك البثّ (على الوسيط) ─────────────────────────
_HERE = os.path.dirname(os.path.abspath(__file__))
_SESS = None          # onnxruntime session (يُحمَّل بخيط خلفيّ عند الإقلاع)
_MC = None            # ثوابت mel
_VOCAB = None
_WINDOW = _MELFB = None
_CMVN = {}
BLANK, NL, DM, LC, TC, CH = 1024, 17, 512, 70, 8, 112
WIN = HOP = NFFT = NR = NMEL = 0
CS = ADV = 0

def _b64f32(s):
    return np.frombuffer(base64.b64decode(s), dtype="<f4").copy()

def _load_model():
    global _SESS, _MC, _VOCAB, _WINDOW, _MELFB, _CMVN, WIN, HOP, NFFT, NR, NMEL, CS, ADV
    try:
        import onnxruntime as ort
        _MC = json.load(io.open(os.path.join(_HERE, "mel_const.json"), encoding="utf-8"))
        _VOCAB = json.load(io.open(os.path.join(_HERE, "vocab.json"), encoding="utf-8"))
        WIN, HOP, NFFT, NR, NMEL = _MC["win_n"], _MC["hop_n"], _MC["n_fft"], _MC["n_rfft"], _MC["n_mels"]
        CS = WIN + (CH - 1) * HOP
        ADV = CH * HOP
        _WINDOW = _b64f32(_MC["window_b64"])
        _MELFB = _b64f32(_MC["mel_fb_b64"]).reshape(NMEL, NR)
        for kind in ("clean", "tlog"):
            _CMVN[kind] = (_b64f32(_MC[f"{kind}_mean_b64"])[:, None], _b64f32(_MC[f"{kind}_std_b64"])[:, None])
        so = ort.SessionOptions()
        so.intra_op_num_threads = max(1, (os.cpu_count() or 2) - 1)
        _SESS = ort.InferenceSession(os.path.join(_HERE, "model.q8.onnx"), so, providers=["CPUExecutionProvider"])
        print("stream model ready", flush=True)
    except Exception as e:
        print("stream model load FAILED:", e, flush=True)

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

def _decode(ids):
    return "".join((_VOCAB[i] if 0 <= i < len(_VOCAB) else "") for i in ids).replace("▁", " ").strip()

class _Stream:
    def __init__(self, kind):
        self.kind = kind if kind in ("clean", "tlog") else "tlog"
        self.buf = np.zeros(0, np.float32)
        self.prev_raw = 0.0
        self.clc = np.zeros((1, NL, LC, DM), np.float32)
        self.clt = np.zeros((1, NL, DM, TC), np.float32)
        self.cll = np.zeros((1,), np.int64)
        self.ids, self.born, self.prev, self.n = [], [], -1, 0

    def feed(self, raw: bytes):
        self.buf = np.concatenate([self.buf, np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0])

    def ready(self):
        return len(self.buf) >= CS

    def _mel(self, seg):
        x = np.empty_like(seg)
        x[0] = seg[0] - 0.97 * self.prev_raw
        x[1:] = seg[1:] - 0.97 * seg[:-1]
        nf = 1 + (len(x) - WIN) // HOP
        from numpy.lib.stride_tricks import as_strided
        fr = as_strided(x, shape=(nf, WIN), strides=(x.strides[0] * HOP, x.strides[0])).copy() * _WINDOW
        sp = np.fft.rfft(fr, n=NFFT, axis=1)
        pw = np.real(sp * np.conj(sp)).astype(np.float32)
        mel = np.log(pw @ _MELFB.T + 2 ** -24).T
        gm, gs = _CMVN[self.kind]
        return ((mel - gm) / (gs + 1e-5)).astype(np.float32)

    def step(self):
        seg = self.buf[:CS]
        mel = self._mel(seg)[None]
        self.prev_raw = float(self.buf[ADV - 1])
        self.buf = self.buf[ADV:]
        o = _SESS.run(None, {"audio_signal": mel, "length": np.array([mel.shape[2]], np.int64),
                             "cache_last_channel": self.clc, "cache_last_time": self.clt,
                             "cache_last_channel_len": self.cll})
        logp, self.clc, self.clt, self.cll = o[0], o[2], o[3], o[4]
        for t in logp[0].argmax(-1):
            t = int(t)
            if t != self.prev and t != BLANK:
                self.ids.append(t)
                self.born.append(self.n)
            self.prev = t
        self.n += 1
        cut = self.n - 4
        recent = [i for i, b in zip(self.ids, self.born) if b >= cut]
        return {"t": self.n, "text": _decode(self.ids), "recent": _decode(recent)}

    def flush(self):
        if len(self.buf) > HOP:
            self.buf = np.concatenate([self.buf, np.zeros(CS - len(self.buf), np.float32)])
            return self.step()
        return None

@app.websocket("/stream")
async def stream(ws: WebSocket):
    ip = _client_ip(ws.scope.get("headers", []), ws.client)
    if not _allow(ip, "ws-dial", 20):
        await ws.close(code=1013)
        return
    with _LOCK:
        if _WS_LIVE[ip] >= 3:
            await ws.close(code=1013)
            return
        _WS_LIVE[ip] += 1
    await ws.accept()
    try:
        # النموذج قد يكون ما زال يُحمَّل (إقلاع بارد) — انتظر حتى 25ث
        waited = 0.0
        while _SESS is None and waited < 25:
            await asyncio.sleep(0.5)
            waited += 0.5
        if _SESS is None:
            await ws.close(code=1011)
            return
        st = _Stream(ws.query_params.get("kind", "tlog"))
        while True:
            msg = await ws.receive()
            if msg.get("type") == "websocket.disconnect":
                break
            if msg.get("bytes") is not None:
                st.feed(msg["bytes"])
                # سقف مخزون: 120ث صوت غير معالَج = إساءة/تسريب → أغلق
                if len(st.buf) > 16000 * 120:
                    await ws.close(code=1009)
                    break
                while st.ready():
                    r = await asyncio.to_thread(st.step)
                    await ws.send_text(json.dumps(r, ensure_ascii=False))
            elif msg.get("text") == "end":
                r = await asyncio.to_thread(st.flush)
                if r:
                    await ws.send_text(json.dumps(r, ensure_ascii=False))
                await ws.close()
                break
    except WebSocketDisconnect:
        pass
    finally:
        with _LOCK:
            _WS_LIVE[ip] = max(0, _WS_LIVE[ip] - 1)

# ───────────────────────── الصحّة + البروكسي (كما كان + حدود معدّل) ─────────────────────────
@app.get("/")
def health():
    return {"ok": True, "muaalem": bool(EP and TOK), "stt": bool(STT_EP and STT_TOK),
            "stream": _SESS is not None}

@app.get("/model")
def model_file(req: Request):
    """نموذج البثّ (126MB) لميزة «على الجهاز» في المتصفّح — Cloudflare Pages يرفض ملفّات >25MiB
    فنقدّمه من هنا (النطاق مسموح في CSP التطبيقين أصلًا). كاش سنة + حدّ 5 تنزيلات/ساعة/IP."""
    if not _allow(_ip_http(req), "model", 5, per=3600):
        return JSONResponse({"error": "rate_limited"}, status_code=429)
    p = os.path.join(_HERE, "model.q8.onnx")
    if not os.path.exists(p):
        return JSONResponse({"error": "no_model"}, status_code=404)
    return FileResponse(p, media_type="application/octet-stream",
                        headers={"Cache-Control": "public, max-age=31536000, immutable"})

@app.post("/stt")
async def stt(req: Request):
    if not _allow(_ip_http(req), "stt", 120):
        return JSONResponse({"error": "rate_limited"}, status_code=429)
    try:
        body = await req.json()
    except Exception:
        return JSONResponse({"error": "invalid_json"}, status_code=400)
    pcm = body.get("pcm")
    if not pcm or not isinstance(pcm, str):
        return JSONResponse({"error": "missing_pcm"}, status_code=400)
    if not (STT_EP and STT_TOK):
        return JSONResponse({"error": "stt_not_configured"}, status_code=503)
    ref = body.get("ref")
    payload = {"inputs": pcm, "ref": ref} if ref else {"inputs": pcm}
    headers = {"Authorization": f"Bearer {STT_TOK}", "Content-Type": "application/json"}
    for _ in range(30):
        try:
            r = requests.post(STT_EP, json=payload, headers=headers, timeout=180)
        except Exception as e:
            return JSONResponse({"error": "stt_error", "detail": str(e)[:200]}, status_code=502)
        if r.status_code in (502, 503):
            time.sleep(6)
            continue
        try:
            data = r.json()
        except Exception:
            return JSONResponse({"text": "", "words": None}, status_code=200)
        obj = data[0] if isinstance(data, list) and data else (data if isinstance(data, dict) else {})
        return JSONResponse({"text": obj.get("text", ""), "words": obj.get("words"),
                             "acoustic": obj.get("acoustic")}, status_code=200)
    return JSONResponse({"text": "", "words": None}, status_code=200)

@app.post("/analyze")
async def analyze(req: Request):
    if not _allow(_ip_http(req), "analyze", 40):
        return JSONResponse({"error": "rate_limited"}, status_code=429)
    try:
        body = await req.json()
    except Exception:
        return JSONResponse({"error": "invalid_json"}, status_code=400)
    pcm = body.get("pcm")
    if not pcm or not isinstance(pcm, str):
        return JSONResponse({"error": "missing_pcm"}, status_code=400)
    if len(pcm) > 30 * 1024 * 1024:
        return JSONResponse({"error": "too_large"}, status_code=413)
    inputs = {"pcm": pcm, "surah": body.get("surah")}
    if body.get("ayahs"):
        inputs["ayahs"] = body["ayahs"]
    else:
        inputs["ayah"] = body.get("ayah")
    if not (EP and TOK):
        return JSONResponse({"error": "not_configured"}, status_code=503)
    headers = {"Authorization": f"Bearer {TOK}", "Content-Type": "application/json"}
    for _ in range(30):  # يعالج البدء البارد للخادم (scale-to-zero → 502/503)
        try:
            r = requests.post(EP, json={"inputs": inputs}, headers=headers, timeout=180)
        except Exception as e:
            return JSONResponse({"error": "proxy_error", "detail": str(e)[:200]}, status_code=502)
        if r.status_code in (502, 503):
            time.sleep(6)
            continue
        try:
            data = r.json()
        except Exception:
            data = {"error": "bad_upstream", "status": r.status_code}
        return JSONResponse(data, status_code=200)
    return JSONResponse({"error": "timeout"}, status_code=200)