# -*- coding: utf-8 -*-
"""معالج HuggingFace Inference Endpoint الكامل لنموذج المُعلِّم (QPS) — يَخلُف المعالج الفونيميّ.

يخلُف `muaalem_handler_deployed.py` الذي يحسب لوجيتات الرؤوس الأحد عشر ثمّ **يفكّ رأس الفونيم
وحده ويرمي الصفات العشر**. هذا المعالج يفكّ الرؤوس الأحد عشر جميعًا على `outs` نفسها (تمريرةٌ
أماميّةٌ واحدة، لا تحميل نموذجٍ ثانٍ) عبر دوالّ `quran_muaalem/decode.py`:

    phonemes_level_greedy_decode  →  chunck_phonemes  →  multilevel_greedy_decode  →  format_sifat

فيُرجِع لكلّ فونيم/كلمةٍ صفاتِها التجويديّة العشر (hams/jahr، shidda/rakhawa، tafkheem/tarqeeq،
itbaq، safeer، qalqla، tikraar، tafashie، istitala، ghonna) مع احتمال كلٍّ، ويقارنها بصفات
المرجع (quran_phonetizer.sifat) فيُطلق تنبيهًا كهرمانيًّا عند اختلاف الصفة بثقةٍ عالية.

توافقٌ رجعيّ كامل: الحقول القديمة (word_hard/word_soft/flagged/ops/match/basmala_*) تبقى كما هي
(محسوبةً من qps.compare/compare_multi على الفونيم المسموع)، والصفات مُضافةٌ فوقها لا كاسرةٌ لها.

ملاحظة على decode.py: دالّة `multilevel_greedy_decode` في الحزمة تحوي خطأً في فرع عدم تطابق
الطول (تُسنِد المُعرِّفات مكان الاحتمالات: `new_probs = ref_aligned_ids[phonemes_mask]`)، فتخرج
احتمالاتٌ = مُعرِّفات (2.0، 3.0…) لبعض الرؤوس. أعدنا تنفيذ الدالّة هنا (`_multilevel_decode`)
بالإصلاح الوحيد (`probs = new_probs[phonemes_mask]`) مع إبقاء بقيّة المنطق حرفيًّا، فتصير كلّ
الاحتمالات ضمن [0,1] وتصلح عتبة التنبيه. باقي الفكّ يستدعي دوالّ الحزمة كما هي.

الطلب:
  {"inputs": {"pcm": "<base64 PCM16 16k أحاديّ>", "surah": 1, "ayahs": [1,2,3]}}
  (أو "ayah": 3 لآيةٍ مفردة؛ "sifat": false لتعطيل الصفات؛ "prob_thresh": 0.6 لضبط العتبة)
المخرج (متوافقٌ رجعيًّا + sifat لكلّ كلمة):
  {"uthmani","words":[{"w","hard","soft","sifat":[{"ph", <10 صفات> }]}],
   "word_hard","word_soft","flagged","ops","match","per","sifat_flagged","timing_ms",...}
"""
import base64
import os
import sys
import time
from typing import Any, Dict

import numpy as np
import torch

MODEL_REPO = os.environ.get("MUAALEM_REPO", "obadx/muaalem-model-v3_2")


class EndpointHandler:
    def __init__(self, path: str = "."):
        sys.path.insert(0, path)  # ليجد qps (muaalem_qps_deployed) داخل مستودع النقطة
        from transformers import AutoFeatureExtractor
        from quran_transcript import QuranPhoneticScriptOutput, chunck_phonemes
        from quran_muaalem.modeling.modeling_multi_level_ctc import (
            Wav2Vec2BertForMultilevelCTC,
        )
        from quran_muaalem.modeling.multi_level_tokenizer import MultiLevelTokenizer
        from quran_muaalem.modeling.vocab import PAD_TOKEN_IDX
        from quran_muaalem.decode import (
            ctc_decode,
            phonemes_level_greedy_decode,
            align_chunked_phonemes_sequence,
            align_predicted_sequence,
        )
        from quran_muaalem.inference import format_sifat
        from quran_muaalem.muaalem_typing import Unit
        import qps

        self.qps = qps
        self.QuranPhoneticScriptOutput = QuranPhoneticScriptOutput
        self._chunck = chunck_phonemes
        self._ctc_decode = ctc_decode
        self._phon_decode = phonemes_level_greedy_decode
        self._align_chunk = align_chunked_phonemes_sequence
        self._align_pred = align_predicted_sequence
        self._format_sifat = format_sifat
        self._Unit = Unit
        self._PAD = PAD_TOKEN_IDX
        # GOP للحركات القصيرة (اختياريّ، غير كاسر): يقيس احتمال الحركة المتوقّعة صوتيًّا بمحاذاةٍ
        # قسريّة بدل word_soft المتذبذب (عتبة مُعايَرة على صوت المستخدم في العميل). تراجعٌ آمن إن غاب.
        try:
            from harakat_gop import harakat_gop_per_word
            self._harakat_gop = harakat_gop_per_word
        except Exception:
            self._harakat_gop = None

        self.model = Wav2Vec2BertForMultilevelCTC.from_pretrained(MODEL_REPO)
        self.tok = MultiLevelTokenizer(MODEL_REPO)
        # طبقة المدّ بالتوقيت (forced alignment) — اختياريّة، تراجعٌ آمن إن غاب torchaudio
        try:
            import madd_timing
            madd_timing.attach(self.tok)
            self._madd = madd_timing
        except Exception:
            self._madd = None
        self.proc = AutoFeatureExtractor.from_pretrained(MODEL_REPO)
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.dtype = torch.float16 if self.device == "cuda" else torch.float32
        self.model.to(self.device, self.dtype).eval()

    # --------------------------- فكّ الرؤوس ---------------------------
    def _multilevel_decode(self, level_to_probs, id_to_vocab, level_to_ref_ids,
                           chunked_pred, ref_chunked, phonemes_units,
                           missing=-100):
        """نسخةٌ مُصحَّحةٌ من quran_muaalem.decode.multilevel_greedy_decode: تفكّ رؤوس الصفات
        وتحاذيها بالفونيم المُتنبَّأ، مع إصلاح إسناد الاحتمالات (السطر المعلَّم FIX)."""
        pad = self._PAD
        level_to_units = {}
        for level in level_to_probs:
            if level == "phonemes":
                continue
            batch_probs, batch_ids = level_to_probs[level].topk(1, dim=-1)
            decode_outs = self._ctc_decode(
                batch_ids.squeeze(-1), batch_probs.squeeze(-1),
                collapse_consecutive=True,
            )
            level_to_units[level] = []
            for seq_idx, dout in enumerate(decode_outs):
                phonemes_mask = torch.BoolTensor(self._align_chunk(
                    ref=ref_chunked[seq_idx], predicted=chunked_pred[seq_idx]))
                if len(dout.ids) != len(chunked_pred[seq_idx]) and (
                    len(chunked_pred[seq_idx]) <= len(ref_chunked[seq_idx])
                ):
                    ref_aligned_ids, mask = self._align_pred(
                        level_to_ref_ids[level][seq_idx], dout.ids,
                        missing_placeholder=missing)
                    probs = dout.p
                    ref_aligned_ids = torch.LongTensor(ref_aligned_ids)
                    mask = torch.BoolTensor(mask)
                    new_probs = torch.zeros(len(ref_aligned_ids), dtype=torch.float32)
                    new_probs[ref_aligned_ids != missing] = probs[mask]
                    ref_aligned_ids[ref_aligned_ids == missing] = pad
                    aligned_ids = ref_aligned_ids[phonemes_mask]
                    probs = new_probs[phonemes_mask]  # FIX: كان ref_aligned_ids[phonemes_mask]
                else:
                    aligned_ids = dout.ids
                    probs = dout.p
                text = "".join(id_to_vocab[level][int(i)] for i in aligned_ids)
                level_to_units[level].append(
                    self._Unit(text=text, probs=probs, ids=aligned_ids))
        level_to_units["phonemes"] = phonemes_units
        return level_to_units

    def _run(self, wave: np.ndarray, ref_obj):
        """تمريرةٌ أماميّةٌ واحدة → (نصّ الفونيم المسموع، صفات المسموع لكلّ مجموعة).
        تفكّ الرؤوس الأحد عشر على outs نفسها عبر دوالّ decode.py (لا نموذج ثانٍ)."""
        feats = self.proc([wave], sampling_rate=16000, return_tensors="pt")
        feats = {k: v.to(self.device, self.dtype) for k, v in feats.items()}
        with torch.no_grad():
            outs = self.model(**feats, return_dict=False)[0]  # dict: 11 رؤوس -> لوجيتات
        probs = {lv: torch.nn.functional.softmax(outs[lv], dim=-1).cpu().to(torch.float32)
                 for lv in outs}
        phonemes_units = self._phon_decode(
            probs["phonemes"], self.tok.id_to_vocab["phonemes"])
        chunked_pred = [self._chunck(u.text) for u in phonemes_units]
        ref_ids = self.tok.tokenize(
            [ref_obj.phonemes], [ref_obj.sifat],
            to_dict=True, return_tensors="pt", padding="longest")["input_ids"]
        ref_chunked = [[s.phonemes for s in ref_obj.sifat]]
        level_to_units = self._multilevel_decode(
            probs, self.tok.id_to_vocab, ref_ids, chunked_pred, ref_chunked,
            phonemes_units)
        sifat = self._format_sifat(level_to_units, chunked_pred, self.tok)[0]
        return phonemes_units[0].text, sifat, outs["phonemes"][0]

    def _ref_obj(self, ref):
        """يبني QuranPhoneticScriptOutput مرجعيًّا للنموذج من مُخرَج qps.build_ref* (بلا فراغ)."""
        return self.QuranPhoneticScriptOutput(
            phonemes=ref["phonemes"].replace(" ", ""),
            sifat=ref["sifat"], mappings=[])

    def _phon_blank_id(self):
        """يكتشف blank رأس الفونيمات: رمزٌ فارغ/<blank>/<pad> في المفردات، وإلّا PAD_TOKEN_IDX، وإلّا الأخير.
        (مُتحقَّق على نموذج obadx/muaalem-model-v3_2: blank=0=[PAD].)"""
        v2i = {v: k for k, v in self.tok.id_to_vocab["phonemes"].items()}
        for marker in ("", "<blank>", "<pad>", "<blk>", "|", "_"):
            if marker in v2i:
                return int(v2i[marker])
        if self._PAD is not None:
            return int(self._PAD)
        return len(self.tok.id_to_vocab["phonemes"]) - 1

    # --------------------------- بناء الرد ---------------------------
    def _result(self, r, pred):
        """الحقول القديمة كما في المعالج المنشور (توافقٌ رجعيّ)."""
        out = {
            "uthmani": r["uthmani"], "match": r["match"], "per": r.get("per"),
            "words": [{"w": w, "hard": r["word_hard"][k] > 0, "soft": r["word_soft"][k]}
                      for k, w in enumerate(r["words"])],
            "word_hard": r["word_hard"], "word_soft": r["word_soft"],
            "word_errors": r["word_hard"], "flagged": r["flagged"],
            "flagged_soft": r["flagged_soft"], "ops": r["ops"],
            "ref_phonemes": r["ref_phonemes"], "pred_phonemes": pred,
        }
        for k in ("basmala_hard", "basmala_soft", "basmala_flagged", "basmala_words"):
            if k in r:
                out[k] = r[k]
        return out

    def _attach_sifat(self, out, ref, pred_sifat, prob_thresh):
        """يقارن الصفات (compare_sifat) ويُرفِق مجموعات الصفات بكلّ كلمةٍ + ملخّص المُنبَّهة."""
        groups = self.qps.compare_sifat(
            ref["sifat"], pred_sifat, ref["sifa_word"], prob_thresh=prob_thresh)
        nw = len(out["words"])
        for wd in out["words"]:
            wd["sifat"] = []
        for g in groups:
            wi = g["word"]
            entry = {k: v for k, v in g.items() if k != "word"}
            if 0 <= wi < nw:
                out["words"][wi]["sifat"].append(entry)
        # ملخّص: الكلمات التي فيها صفةٌ مُنبَّهة + أنواع الصفات المُنبَّهة لكلّ كلمة.
        flagged, word_sifat = [], [0] * nw
        for g in groups:
            wi = g["word"]
            if g["mismatch"] and 0 <= wi < nw:
                word_sifat[wi] = 1
                mis = [lvl for lvl in self.qps.SIFAT_LEVELS if g[lvl]["mismatch"]]
                flagged.append({"word": ref["words"][wi] if wi < len(ref["words"]) else "",
                                "index": wi, "ph": g["ph"], "heard_ph": g["heard_ph"],
                                "sifat": mis})
        out["word_sifat"] = word_sifat
        out["sifat_flagged"] = flagged
        out["prob_thresh"] = prob_thresh
        return out

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        t0 = time.time()
        inp = data.get("inputs", data)
        if isinstance(inp, str):
            inp = {"pcm": inp, **{k: v for k, v in data.items() if k != "inputs"}}
        try:
            raw = base64.b64decode(inp["pcm"])
            x = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
            if x.size == 0:
                return {"error": "empty audio"}
            pk = float(np.max(np.abs(x))) or 1.0
            x = (x / pk) * 0.97
            surah = int(inp["surah"])
            ayahs = [int(a) for a in inp["ayahs"]] if inp.get("ayahs") else None
            want_sifat = inp.get("sifat", True)
            prob_thresh = float(inp.get("prob_thresh", self.qps.SIFAT_PROB_THRESH))

            # المرجع (يحوي الآن sifat وsifa_word) — تمريرةٌ واحدة على كامل الصوت (لا تجزئة):
            # الفكّ الكامل يحتاج مرجعًا واحدًا متّسقًا لمحاذاة الصفات؛ زوائد المنطوق (بسملة/ضوضاء)
            # تُعالَج بمحاذاة المجموعات (بلا نظير) داخل compare_sifat، وبقصّ البسملة في qps للفونيم.
            if ayahs:
                ref = self.qps.build_ref_multi(surah, ayahs)
            else:
                ref = self.qps.build_ref(surah, int(inp["ayah"]))

            t_infer = time.time()
            pred_ph, pred_sifat, ph_logits = self._run(x, self._ref_obj(ref))
            infer_ms = int((time.time() - t_infer) * 1000)

            # الحقول القديمة (فونيم فقط) عبر منطق qps المُتحقَّق (قصّ البسملة/الحروف المقطّعة…).
            if ayahs:
                r = self.qps.compare_multi(surah, ayahs, pred_ph)
            else:
                r = self.qps.compare(surah, int(inp["ayah"]), pred_ph)
            out = self._result(r, pred_ph)

            if want_sifat:
                out = self._attach_sifat(out, ref, pred_sifat, prob_thresh)
            out["mode"] = "whole-sifat" if want_sifat else "whole"
            if self._madd is not None:
                try:
                    out["madd"] = self._madd.madd_flags(
                        ph_logits, surah, ayahs or [int(inp["ayah"])], ref)
                except Exception:
                    pass
            # GOP للحركات القصيرة: إشارةٌ صوتيّة مستقرّة (احتمال الحركة المتوقّعة مقابل الأعلى) بمحاذاةٍ
            # قسريّة — تُغني العميل عن word_soft المتذبذب. إضافةٌ غير كاسرة (حقولٌ إضافيّة فقط، try/except).
            # الفهارس بكلمات المرجع الفونيميّ (قد تختلف عن word_soft عند الدمج)؛ العميل يحاذيها بـalignRefToMu.
            if self._harakat_gop is not None:
                try:
                    pp = torch.nn.functional.softmax(
                        ph_logits.to(torch.float32), dim=-1).cpu().numpy()   # (T,V) احتمالات رأس الفونيمات
                    lp = np.log(pp + 1e-9)
                    hg = self._harakat_gop(lp, ref["phonemes"], self.tok.id_to_vocab["phonemes"],
                                           blank_id=self._phon_blank_id())
                    out["harakat_gop"] = hg["per_word"]
                    out["harakat_flagged"] = hg["flagged"]
                    out["harakat_ref_words"] = ref["phonemes"].split(" ")
                except Exception as _e:
                    out["harakat_gop_error"] = str(_e)[:200]
            out["timing_ms"] = {"infer": infer_ms, "total": int((time.time() - t0) * 1000)}
            return out
        except Exception as e:
            import traceback
            # لا نُعيد أثر الاستثناء (مسارات ملفّاتٍ داخليّة) للمتصل — نطبعه في سجلّ الخادم فقط
            # ونُرجع رسالة عامّة نظيفة. (اكتُشِف هذا في تدقيقٍ أمنيّ: طلبٌ فاسدٌ/خارج المدى كان
            # يُرجع تتبّعًا كاملًا يكشف بنية الملفّات الداخليّة.)
            print(f"[al-moratel-qps] خطأ: {e}\n{traceback.format_exc()}", flush=True)
            return {"error": "processing_failed"}