"""معالِج الطبقة الصوتية الكاملة (Tarteel-class) — يتبنّى حزمة المؤلّف tajweed/.

لكل كلمة من النصّ المرجعيّ يُرجع:
  • توقيت (start_ms / end_ms) من المحاذاة القسرية CTC،
  • prob = احتمال صحّة النطق (رأس النطق المُدرَّب 1.33M معامل)،
  • gop = جودة النطق (Goodness of Pronunciation، Witt & Young)،
  • status = correct / warning / wrong (دمج الرأس + GOP).

بروتوكول الطلب:
  {"inputs": "<pcm16 b64>", "ref": "<نصّ الكلمات بالتشكيل>"}  →
  {"text": "...", "words": [{w, start_ms, end_ms, prob, gop, status}], "provider": "quran-stt-acoustic"}

تدهور آمن: إن غابت تبعية/نموذج، يعمل بالنصّ فقط (لا ينهار). متوافق مع البروتوكول القديم
({"inputs": pcm} بلا ref → نصّ فقط).
"""
from __future__ import annotations

import base64
import os
import time
import traceback
from typing import Any, Dict, List, Optional

# Tame OpenMP/BLAS pools BEFORE torch/numpy load their runtimes. On a CPU-only HF
# image the base torch may be a CUDA build whose Intel-OpenMP pool spins across every
# host core; combined with onnxruntime's own pool this oversubscribes shared nodes and
# makes inference ~40x slower. The ONNX pool is bounded in tajweed/aligner.py; these
# env caps bound torch/BLAS the same way. QURAN_STT_THREADS overrides (see aligner).
_omp = os.environ.get("QURAN_STT_THREADS", "").strip()
_omp = _omp if (_omp.isdigit() and int(_omp) > 0) else "4"
for _v in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS"):
    os.environ.setdefault(_v, _omp)

import numpy as np

# DIAG: اطبع نسخ المكتبات مرّةً عند الاستيراد لكشف انحدار numpy/ort على الصورة الأحدث.
try:
    import onnxruntime as _ort_diag
    _ort_v = _ort_diag.__version__
except Exception as _e:
    _ort_v = f"ERR:{_e}"
try:
    import torch as _torch_diag
    _torch_v = _torch_diag.__version__
except Exception as _e:
    _torch_v = f"ERR:{_e}"
print(f"DIAG versions: numpy={np.__version__} onnxruntime={_ort_v} torch={_torch_v} "
      f"OMP={os.environ.get('OMP_NUM_THREADS')} cpu_count={os.cpu_count()}", flush=True)

_HERE = os.path.dirname(os.path.abspath(__file__))


def _decode_pcm(b64: str) -> np.ndarray:
    raw = base64.b64decode(b64)
    return np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0


# تطبيع الرسم العثمانيّ إلى فضاء النموذج (وإلا تصير ٱ/ٰ رموز <unk> فتُحسَب خطأ).
# النموذج يُخرج ا كاملة لألف الوصل وألف الخنجريّة (madd) — فنحوّلها لا نحذفها.
_NORM_MAP = {
    0x0671: "ا",  # ٱ ألف وصل → ا
    0x0623: "ا",  # أ → ا
    0x0625: "ا",  # إ → ا
    0x0622: "ا",  # آ → ا
    0x0670: "ا",  # ٰ ألف خنجرية (مدّ) → ا
}
# علامات قرآنية صغيرة تُزال (لا تؤثّر في النطق ضمن فضاء النموذج).
_STRIP = set(range(0x06D6, 0x06EE)) | {0x0640, 0x0653, 0x0654, 0x0655, 0x0656}


def _norm_ref(s: str) -> str:
    out = []
    for ch in s:
        o = ord(ch)
        if o in _STRIP:
            continue
        out.append(_NORM_MAP.get(o, ch))
    return "".join(out)


# دمج إشارة الرأس + GOP إلى حالة كلمة واحدة.
def _word_status(prob: Optional[float], gop: Optional[float]) -> str:
    # الرأس المُدرَّب هو الإشارة الأساسية (تمييزيّ، مُعايَر).
    if prob is not None:
        if prob >= 0.40:
            base = "correct"
        elif prob >= 0.30:
            base = "warning"
        else:
            base = "wrong"
    else:
        base = None
    # GOP إشارة مكمّلة (سالبة؛ قرب الصفر = أصحّ).
    if gop is not None:
        if gop < -2.0:
            g = "wrong"
        elif gop < -0.5:
            g = "warning"
        else:
            g = "correct"
    else:
        g = None
    # ادمج: خذ الأسوأ بين الإشارتين (احترازيّ — لا نُمرّر خطأً صامتًا).
    order = {"correct": 0, "warning": 1, "wrong": 2}
    cands = [x for x in (base, g) if x is not None]
    if not cands:
        return "correct"
    return max(cands, key=lambda s: order[s])


class EndpointHandler:
    def __init__(self, path: str = _HERE):
        self.aligner = None
        self.head = None
        self.gop_fn = None
        self.tokenizer = None
        self.init_error = None
        try:
            # كشف الجهاز: CUDA إن توفّر (نقطة GPU)، وإلا CPU — تراجع آمن تامّ فلا تنهار
            # الخدمة إن غاب CUDA. جلسة ONNX في المحاذي تكتشف CUDAExecutionProvider ذاتيًّا.
            try:
                import torch as _torch
                _device = "cuda" if _torch.cuda.is_available() else "cpu"
            except Exception:
                _device = "cpu"
            print(f"  device: {_device}")
            # CTCAligner يحمّل model_with_encoder (logprobs + encoder_output) + tokenizer.
            from tajweed.aligner import CTCAligner, ctc_forced_align, TokenInterval
            from tajweed.gop_scorer import gop_from_logprobs
            self._CTCAligner = CTCAligner
            self._ctc_forced_align = ctc_forced_align
            self._TokenInterval = TokenInterval
            self.gop_fn = gop_from_logprobs
            model = os.path.join(path, "model_with_encoder.q8.ort")
            tok = os.path.join(path, "tokenizer.model")
            # log-mel على GPU إن توفّر (أسرع)، وإلا CPU.
            self.aligner = CTCAligner(model, tok, cuda_feats=(_device == "cuda"))
            self.tokenizer = self.aligner.tokenizer
            if not self.aligner._has_encoder_output():
                # نموذج بلا encoder_output → الرأس معطّل، لكن المحاذاة+GOP تعملان.
                print("WARN: ONNX بلا encoder_output؛ رأس النطق معطّل.")
            else:
                head_ckpt = os.path.join(path, "head", "pronunciation_head.pt")
                if os.path.exists(head_ckpt):
                    try:
                        from tajweed.head_scorer import HeadPronunciationScorer
                        self.head = HeadPronunciationScorer(head_ckpt, device=_device)
                    except Exception as e:
                        print(f"WARN: فشل تحميل رأس النطق ({e})؛ يُكتفى بـ GOP.")
                        self.head = None
        except Exception as e:
            # تدهور آمن تامّ: نصّ فقط عبر معالِج المحاذاة الخفيف.
            self.init_error = str(e)
            print(f"WARN: تعذّر تهيئة الطبقة الصوتية الكاملة ({e}).")

        # محرّك التجويد الـDSP (Tier-1): قياس صوتيّ فعليّ لأحكام المدّ/الغنّة/القلقلة/
        # التفخيم/الإخفاء... لكل حرف. تدهور آمن مستقلّ: إن غاب لا يُعطّل المحاذاة+الرأس+GOP.
        self.engine = None
        self.analyze_text = None
        self._ARABIC_LETTERS = None
        if self.aligner is not None:
            try:
                from tajweed.engine import TajweedEngine
                from tajweed.text_analyzer import analyze_text, ARABIC_LETTERS
                self.engine = TajweedEngine(sample_rate=16000)
                self.analyze_text = analyze_text
                self._ARABIC_LETTERS = ARABIC_LETTERS
            except Exception as e:
                print(f"WARN: تعذّر تحميل محرّك التجويد الـDSP ({e})؛ يُكتفى بالنطق/GOP.")
                self.engine = None

    # ---- نصّ فقط (فكّ جشِع) عبر جلسة المحاذي ----
    def _greedy_text(self, logprobs: np.ndarray) -> str:
        ids = logprobs.argmax(axis=-1)
        out, prev = [], -1
        BLANK = self.aligner.BLANK_ID
        for i in ids:
            i = int(i)
            if i != prev and i != BLANK and 0 <= i < self.tokenizer.get_piece_size():
                out.append(self.tokenizer.id_to_piece(i))
            prev = i
        return "".join(out).replace("▁", " ").strip()

    def _align_to_words(self, logprobs, enc, audio, ref: str) -> List[Dict[str, Any]]:
        HOP = self.aligner.OUTPUT_HOP_S
        ref_norm = _norm_ref(ref)
        token_ids = self.tokenizer.encode(ref_norm, out_type=int)
        if not token_ids:
            return []
        intervals = self._ctc_forced_align(logprobs, token_ids, blank_id=self.aligner.BLANK_ID)
        token_intervals = [
            self._TokenInterval(
                token_id=tid, token_str=self.tokenizer.id_to_piece(tid),
                start_s=a * HOP, end_s=b * HOP,
            )
            for tid, (a, b) in zip(token_ids, intervals)
        ]
        # درجات الرأس لكل رمز (إن توفّر).
        head_by_idx: Dict[int, Any] = {}
        if self.head is not None and enc is not None:
            try:
                for i, hs in enumerate(self.head.score(enc, token_intervals, output_hop_s=HOP)):
                    head_by_idx[i] = hs
            except Exception as e:
                print(f"WARN: فشل تسجيل الرأس ({e}).")
        # GOP لكل رمز.
        gop_by_idx: Dict[int, float] = {}
        if self.gop_fn is not None:
            try:
                for i, gr in enumerate(self.gop_fn(logprobs, token_intervals, output_hop_s=HOP,
                                                   blank_id=self.aligner.BLANK_ID)):
                    gop_by_idx[i] = gr.gop_normalized
            except Exception as e:
                print(f"WARN: فشل GOP ({e}).")
        # تجميع الرموز إلى كلمات (الكلمة تبدأ عند رمز يبدأ بـ ▁).
        words: List[Dict[str, Any]] = []
        cur: Optional[Dict[str, Any]] = None
        for i, ti in enumerate(token_intervals):
            piece = ti.token_str
            starts_word = piece.startswith("▁") or cur is None
            if starts_word:
                if cur is not None:
                    words.append(cur)
                cur = {"w": "", "start_s": ti.start_s, "end_s": ti.end_s,
                       "fa": int(intervals[i][0]), "fb": int(intervals[i][1]),
                       "probs": [], "gops": []}
            cur["w"] += piece.replace("▁", "")
            cur["end_s"] = ti.end_s
            cur["fb"] = int(intervals[i][1])  # نطاق إطارات الكلمة (للفكّ الموضعيّ)
            if i in head_by_idx:
                cur["probs"].append(head_by_idx[i].prob_correct)
            if i in gop_by_idx:
                cur["gops"].append(gop_by_idx[i])
        if cur is not None:
            words.append(cur)
        # فكّ جشِع **موضعيّ** على نطاق إطارات كل كلمة — إشارة آمنة لتأكيد أن صوت المقطع يحوي
        # الكلمة المتوقَّعة فعلًا (يكشف ما يُسقطه/يخطئه التفريغ العامّ على الصوت الطويل). يُقارَن
        # بالمتوقَّع في العميل لإنقاذ كلمةٍ عُلِّمت خطأً نصيًّا لكنّها نُطِقت صحيحةً.
        BLANK = self.aligner.BLANK_ID
        def _local_decode(fa: int, fb: int) -> str:
            if fb < fa:
                return ""
            ids = logprobs[fa:fb + 1].argmax(axis=-1)
            o, prev = [], -1
            for x in ids:
                x = int(x)
                if x != prev and x != BLANK and 0 <= x < self.tokenizer.get_piece_size():
                    o.append(self.tokenizer.id_to_piece(x))
                prev = x
            return "".join(o).replace("▁", "")
        # اختزال إلى المخرجات النهائية لكل كلمة.
        out: List[Dict[str, Any]] = []
        for wd in words:
            prob = round(float(min(wd["probs"])), 3) if wd["probs"] else None   # أضعف حرف يحكم
            gop = round(float(np.mean(wd["gops"])), 3) if wd["gops"] else None
            out.append({
                "w": wd["w"],
                "start_ms": int(round(wd["start_s"] * 1000)),
                "end_ms": int(round(wd["end_s"] * 1000)),
                "prob": prob,
                "gop": gop,
                "status": _word_status(prob, gop),
                "decoded": _local_decode(wd.get("fa", 0), wd.get("fb", -1)),
            })
        # إثراء التجويد الـDSP (مدّ/غنّة/قلقلة/تفخيم/إخفاء...) لكل كلمة. تدهور آمن:
        # أيّ فشل يُبقي الكلمات بتوقيت+نطق+GOP دون تجويد.
        if self.engine is not None and self.analyze_text is not None:
            try:
                self._attach_tajweed(audio, token_intervals, out, ref_norm)
            except Exception as e:
                print(f"WARN: فشل إثراء التجويد ({e}).")
        return out

    # ---- إثراء التجويد الـDSP: يربط أحكام كل حرف بكلمتها بالتوقيت المُحاذى ----
    def _attach_tajweed(self, audio, token_intervals, out, ref_norm: str) -> None:
        """يبني توقيتًا لكل حرف من الرموز المُحاذاة (بترجيح حروف المدّ)، ثم يشغّل
        محلّل النصّ + محرّك القياس الصوتيّ، ويُسنِد أحكام كل حرف إلى كلمته بالفهرس.

        النصّ المُطبَّع (ref_norm) نفسه يُستعمل للتحليل والتوكنة → تطابق ١:١ بين
        تسلسل الحروف وتوقيت الرموز (لا انجراف). يُبقي التطبيع الحركات/السكون/الشدّة
        التي تحتاجها معظم القواعد، ويسقط فقط همزة الوصل وعلامات صغيرة.
        """
        MADD_LETTERS = {"ا", "و", "ي", "ى", "آ"}
        MADD_WEIGHT = 2.5  # حروف المدّ تستهلك وقتًا أطول → ترجيح أعلى عند تقسيم الفاصل
        # (١) توقيت كل حرف + الكلمة التي ينتمي إليها (إعادة نفس منطق تجميع الرموز للكلمات).
        char_times: List[tuple] = []
        char_word: List[int] = []
        wi = -1
        for i, ti in enumerate(token_intervals):
            piece = ti.token_str
            if piece.startswith("▁") or i == 0:
                wi += 1
            letters = [c for c in piece.replace("▁", "") if c in self._ARABIC_LETTERS]
            if not letters:
                continue
            weights = [MADD_WEIGHT if c in MADD_LETTERS else 1.0 for c in letters]
            total_w = sum(weights)
            unit = (ti.end_s - ti.start_s) / total_w if total_w else 0.0
            cur = ti.start_s
            for w in weights:
                seg = unit * w
                char_times.append((cur, cur + seg))
                char_word.append(wi)
                cur += seg
        if not char_times:
            return
        # (٢) تحليل النصّ (سياق عابر للكلمات محفوظ) + قياس صوتيّ لكل حرف.
        annotations = self.analyze_text(ref_norm)
        n = min(len(annotations), len(char_times))
        if n == 0:
            return
        results = self.engine.score_with_annotations(audio, annotations[:n], char_times[:n])
        # (٣) إسناد أحكام كل حرف إلى كلمته. نُبقي القواعد القابلة للقياس فقط (flag != n/a).
        per_word: List[List[dict]] = [[] for _ in out]
        for idx, r in enumerate(results):
            wj = char_word[idx]
            if not (0 <= wj < len(out)):
                continue
            for m in r.measurements:
                flag = m.get("flag")
                if not flag or flag == "n/a":
                    continue
                per_word[wj].append({
                    "rule": m.get("rule", ""),
                    "flag": flag,                       # ok | short | long | missing | weak ...
                    "msg_ar": m.get("msg_ar", ""),
                })
        for wd, tj in zip(out, per_word):
            if tj:
                wd["tajweed"] = tj

    # ---- تفريغ مُجزّأ عند الوقفات (VAD) — يحلّ تدهور المودل على الصوت الطويل ----
    def _vad_segments(self, audio: np.ndarray, sr: int = 16000):
        """نقاط تقطيع عند الصمت (>0.35ث): يعطي مقاطع بطول آية يتقنها المودل المدرَّب
        على مقاطع قصيرة. الصوت مُطبَّع (ذروة ~0.7) فالعتبة المطلقة 0.02 مستقرّة."""
        fr = int(sr * 0.03)
        n = (len(audio) - fr) // fr
        if n < 2:
            return [(0, len(audio))]
        rms = np.sqrt(np.array([(audio[i * fr:i * fr + fr] ** 2).mean() for i in range(n)]) + 1e-9)
        sil = rms < 0.02
        pad = int(0.2 * sr)
        segs, in_sp, st, k = [], False, 0, 0
        while k < n:
            if not sil[k]:
                if not in_sp:
                    in_sp, st = True, k
                k += 1
            else:
                j = k
                while j < n and sil[j]:
                    j += 1
                if in_sp and ((j - k) * 0.03 >= 0.35 or j >= n):
                    segs.append((max(0, st * fr), min(len(audio), k * fr + pad)))
                    in_sp = False
                k = j
        if in_sp:
            segs.append((st * fr, len(audio)))
        # حارس: مقطع طويل بلا وقفة (آية طويلة) يُقسَّم نوافذ ~8ث كي لا يتدهور.
        out = []
        for a, b in (segs or [(0, len(audio))]):
            if b - a > 12 * sr:
                w, ov = 8 * sr, 1 * sr
                t = a
                while t < b:
                    out.append((t, min(b, t + w)))
                    t += w - ov
            else:
                out.append((a, b))
        return out

    def _segmented_text(self, audio: np.ndarray, logprobs: np.ndarray) -> str:
        """نصّ مُجزّأ للصوت الطويل (>9ث): يُفرّغ كل مقطع وحده ويجمعها (يستعيد البسملة
        والكلمات التي تُسقطها التمريرة الواحدة). للقصير/اللحظيّ: تمريرة واحدة سريعة."""
        if len(audio) <= 9 * 16000:
            return self._greedy_text(logprobs)
        parts = []
        for a, b in self._vad_segments(audio):
            seg = audio[a:b]
            if len(seg) < int(0.3 * 16000):
                continue
            try:
                lp = (self.aligner._run_model_with_encoder(seg)[0]
                      if self.aligner._has_encoder_output() else self.aligner._run_model(seg))
                t = self._greedy_text(lp)
                if t:
                    parts.append(t)
            except Exception as e:
                print(f"WARN: فشل تفريغ مقطع ({e}).")
        joined = " ".join(parts).strip()
        return joined or self._greedy_text(logprobs)

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        inp = data.get("inputs", data)
        if isinstance(inp, dict):
            pcm_b64 = inp.get("pcm")
            ref = data.get("ref") or inp.get("ref")
        else:
            pcm_b64 = inp
            ref = data.get("ref")
        if not pcm_b64:
            return {"error": "missing_pcm"}
        if self.aligner is None:
            return {"text": "", "words": [], "provider": "quran-stt-acoustic",
                    "error": self.init_error or "uninitialized"}
        try:
            _t = time.time()
            audio = _decode_pcm(pcm_b64)
            print(f"DIAG call: decoded {len(audio)} samples ({len(audio)/16000:.2f}s) in {time.time()-_t:.3f}s", flush=True)
            # تمريرة واحدة: logprobs (+encoder إن توفّر).
            _t = time.time()
            if self.aligner._has_encoder_output():
                logprobs, enc = self.aligner._run_model_with_encoder(audio)
            else:
                logprobs, enc = self.aligner._run_model(audio), None
            print(f"DIAG call: ONNX inference in {time.time()-_t:.3f}s (logprobs {getattr(logprobs,'shape',None)})", flush=True)
            _t = time.time()
            text = self._segmented_text(audio, logprobs)
            print(f"DIAG call: text in {time.time()-_t:.3f}s -> '{text[:40]}'", flush=True)
            words = []
            if ref:
                try:
                    _t = time.time()
                    words = self._align_to_words(logprobs, enc, audio, ref)
                    print(f"DIAG call: align in {time.time()-_t:.3f}s ({len(words)} words)", flush=True)
                except Exception as e:
                    print(f"WARN: فشل المحاذاة الكاملة ({e}).", flush=True)
                    words = []
            return {"text": text, "words": words, "provider": "quran-stt-acoustic"}
        except Exception as e:
            print("DIAG call ERROR:\n" + traceback.format_exc(), flush=True)
            return {"text": "", "words": [], "provider": "quran-stt-acoustic", "error": f"call:{e}"}