"""
term_extractor.py -- LLM + XLMR hate term extraction.
English: Qwen LLM (when cached) or XLMR similarity fallback
Amharic: Afro-XLMR semantic similarity (n-gram + similarity scoring)
Oromo:   Afro-XLMR semantic similarity (n-gram + similarity scoring)
"""

import os, sys, re, json

BASE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, BASE)


def detect_language(text):
    am = sum(1 for c in str(text) if "\u1200" <= c <= "\u137f")
    if am > 2:
        return "amharic"
    oromo_markers = ["dha", "keenya", "biyya", "qabna", "jirti", "isaan", "kana",
                     "qabna", "barbaachisaa", "balleessaa", "hordoftoota"]
    if any(w in str(text).lower().split() for w in oromo_markers):
        return "oromo"
    return "english"


def _get_llm(lang="english"):
    import torch
    hf_token = os.environ.get("HF_TOKEN", None)
    _pn  = sys.modules.get("pattern_namer")
    _reg = getattr(_pn, "_model_registry", {}) if _pn else {}
    if lang == "amharic" and "CohereLabs/aya-expanse-8b" in _reg:
        return _reg["CohereLabs/aya-expanse-8b"]
    model_name = "Qwen/Qwen2.5-1.5B-Instruct"
    if model_name in _reg:
        return _reg[model_name]
    from transformers import AutoTokenizer, AutoModelForCausalLM
    tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForCausalLM.from_pretrained(
        model_name, torch_dtype=torch.bfloat16, device_map="auto", token=hf_token
    )
    model.eval()
    if _pn is not None:
        if not hasattr(_pn, "_model_registry"):
            _pn._model_registry = {}
        _pn._model_registry[model_name] = (tokenizer, model)
    return tokenizer, model


_COMMON_SINGLES = {
    "islam","muslim","muslims","christian","christians","orthodox",
    "jewish","jew","oromo","amhara","tigray","ethiopia","ethiopian",
    "kill","destroy","destroying","hate","disease","culture","women",
    "church","churches","mosque","mosques","religion","people",
}


def _get_xlmr_enc():
    """Get XLMR encoder -- tries ml_engine first, then loads directly."""
    import sys
    _me = sys.modules.get("detector.ml_engine")
    if _me is None:
        try:
            import detector.ml_engine as _me
        except Exception:
            return None
    enc = getattr(_me, "_xlmr_enc", None)
    if enc is None:
        try:
            _me._clf = None
            _me._load()
            enc = getattr(_me, "_xlmr_enc", None)
        except Exception as e:
            print(f"  [xlmr] load failed: {e}")
    return enc


def _is_meaningful_gram(gram, source_text=""):
    """Filter out grammatical fragments that carry no standalone hate signal."""
    words = gram.split()
    # Ethiopic: reject if all words are very short (likely particles/conjunctions)
    am_words = [w for w in words if any("\u1200" <= c <= "\u137f" for c in w)]
    if am_words and all(len(w) <= 3 for w in am_words):
        return False
    # Reject pure connectors (Amharic common particles)
    AM_PARTICLES = {"ነው","ናቸው","ነበር","ይሆናል","ነበሩ","ሲሆን","እና","ወይም",
                    "ግን","ስለዚህ","ነገር","ሆኖም","አለ","አሉ","ነን","ናቸን"}
    if am_words and all(w in AM_PARTICLES for w in am_words):
        return False
    # Latin: reject if all words are stopwords or particles
    lat_words = [w for w in words if all("a" <= c <= "z" for c in w.lower())]
    LAT_PARTICLES = {"the","a","an","is","it","in","on","at","to","for","of",
                     "and","or","but","not","are","was","be","as","by","with",
                     "from","have","has","they","we","dha","fi","irratti"}
    if lat_words and all(w.lower() in LAT_PARTICLES for w in lat_words):
        return False
    return True


def _xlmr_score(ngrams, hate_anchors, threshold=0.75, max_keep=10):
    """Score n-grams against hate anchors using XLMR similarity."""
    import torch
    import torch.nn.functional as F
    import numpy as np

    enc = _get_xlmr_enc()
    if enc is None:
        return []

    def encode(texts):
        arr  = enc._encode_batch(texts)
        embs = torch.tensor(np.array(arr), dtype=torch.float32)
        return F.normalize(embs, dim=-1)

    # Pre-filter nonsense grams before scoring
    filtered_ngrams = {g: v for g, v in ngrams.items()
                       if _is_meaningful_gram(g, v.get("text",""))}
    if not filtered_ngrams:
        return []

    gram_list   = list(filtered_ngrams.keys())[:40]
    anchor_embs = encode(hate_anchors)
    gram_embs   = encode(gram_list)
    sims        = torch.mm(gram_embs, anchor_embs.T)
    best_sims   = sims.max(dim=1).values

    scored = sorted(
        [(gram_list[i], best_sims[i].item()) for i in range(len(gram_list))
         if best_sims[i].item() >= threshold],
        key=lambda x: -x[1]
    )

    kept_terms = []
    results    = []
    for gram, sim in scored:
        g_words = gram.split()
        already_covered = any(
            all(w in k.split() for w in g_words) and len(k.split()) > len(g_words)
            for k in kept_terms
        )
        if already_covered:
            continue
        kept_terms.append(gram)
        results.append((gram, sim))
        if len(results) >= max_keep:
            break
    return results


def _get_lexicon_words():
    """Get known hate words from the lexicon DB for anchor filtering."""
    try:
        from detector.models import LexiconEntry
        words = set()
        for term in LexiconEntry.objects.exclude(
                grouped_label="Normal").values_list("term", flat=True):
            for w in str(term).strip().split():
                if len(w) > 2:
                    words.add(w.lower())
        return words
    except Exception:
        return set()


def _extract_amharic_terms(batch, existing_lower, all_results):
    HATE_ANCHORS_AM = [
        "\u12a0\u123b\u1263\u122a\u12cd\u127d",
        "\u130b\u120b\u1276\u127b\u12c8\u1295",
        "\u12cd\u123b \u1293\u1278\u12c8",
        "\u1218\u1263\u1228\u122d \u12a0\u1208\u1263\u1278\u12c8",
        "\u12ad\u122d\u1235\u1272\u12eb\u1293\u127d\u1295 \u12eb\u1243\u1325\u120b\u120d",
        "\u130d\u1295\u1299\u12cd\u1295 \u1270\u12cd\u1290\u1271",
        "\u1208\u1240\u12c6\u1295 \u1270\u12cd\u1290\u1271",
    ]
    ngrams = {}
    for r in batch:
        text  = r.get("text", "")
        label = r.get("label") or r.get("predicted_label", "Identity-Based Hate")
        words = re.findall(r"[\u1200-\u137f]+", text)
        for n in range(2, 5):
            for i in range(len(words) - n + 1):
                gram = " ".join(words[i:i+n])
                if gram.lower() not in existing_lower and gram not in ngrams:
                    ngrams[gram] = {"text": text, "label": label}
    if not ngrams:
        return
    try:
        # Only keep n-grams containing at least one known hate word from lexicon
        lexicon_words = _get_lexicon_words()
        if lexicon_words:
            filtered = {}
            for gram, meta in ngrams.items():
                gram_words = {w.lower() for w in gram.split()}
                if gram_words & lexicon_words:
                    filtered[gram] = meta
            ngrams = filtered if filtered else ngrams
            print(f"  [amharic_extract] {len(ngrams)} n-grams contain lexicon words")

        scored = _xlmr_score(ngrams, HATE_ANCHORS_AM, threshold=0.78)
        for gram, sim in scored:
            if any(c["term"].lower() == gram.lower() for c in all_results):
                continue
            # Prefer shorter focused phrases over long sentence fragments
            if len(gram.split()) > 4:
                continue
            meta = ngrams[gram]
            all_results.append({
                "term":              gram,
                "suggested_label":   meta["label"],
                "extraction_method": "amharic_xlmr",
                "score":             round(sim, 3),
                "source_text":       meta["text"][:300],
                "language":          "amharic",
            })
            print(f"  [amharic_extract] sim={sim:.3f} Kept: \'{gram}\'")
    except Exception as e:
        print(f"  [amharic_extract] Error: {e}")


def _extract_oromo_terms(batch, existing_lower, all_results):
    HATE_ANCHORS_OM = [
        "diina keenya",
        "ari\'uu qabna",
        "baasuun barbaachisaa",
        "ajjeesuu qabna",
        "balleessuun barbaachisaa",
        "gad aanaa dha",
        "biyya keenya keessaa",
    ]
    EN_STOP = {"the","a","an","is","it","in","on","at","to","for","of","and",
               "or","but","not","are","was","be","as","by","this","that","with",
               "from","have","has","they","we","you","i","he","she","our","their"}
    ngrams = {}
    for r in batch:
        text  = r.get("text", "")
        label = r.get("label") or r.get("predicted_label", "Identity-Based Hate")
        words = [w for w in re.findall(r"[a-zA-Z\']+", text.lower())
                 if len(w) > 3 and w not in EN_STOP]
        for n in range(2, 5):
            for i in range(len(words) - n + 1):
                gram = " ".join(words[i:i+n])
                if gram not in existing_lower and gram not in ngrams:
                    ngrams[gram] = {"text": text, "label": label}
    if not ngrams:
        return
    try:
        scored = _xlmr_score(ngrams, HATE_ANCHORS_OM)
        for gram, sim in scored:
            if any(c["term"].lower() == gram.lower() for c in all_results):
                continue
            meta = ngrams[gram]
            all_results.append({
                "term":              gram,
                "suggested_label":   meta["label"],
                "extraction_method": "oromo_xlmr",
                "score":             round(sim, 3),
                "source_text":       meta["text"][:300],
                "language":          "oromo",
            })
            print(f"  [oromo_extract] sim={sim:.3f} Kept: \'{gram}\'")
    except Exception as e:
        print(f"  [oromo_extract] Error: {e}")


def _extract_english_terms_xlmr(batch, existing_lower, all_results):
    HATE_ANCHORS_EN = [
        "kill them all",
        "must be expelled",
        "should be burned",
        "muslim savages",
        "islamic extremists",
        "jihadists terrorists",
        "subhuman parasites",
        "expel them all",
        "infidels must die",
        "ethnic cleansing",
        "cockroaches vermin",
        "must be eliminated",
    ]
    EN_STOP = {"the","a","an","is","it","in","on","at","to","for","of","and",
               "or","but","not","are","was","be","as","by","this","that","with",
               "from","have","has","they","we","you","i","he","she","our","their",
               "these","those","who","what","should","would","could","will","all"}
    ngrams = {}
    for r in batch:
        text  = r.get("text", "")
        label = r.get("label") or r.get("predicted_label", "Identity-Based Hate")
        words = [w.lower() for w in re.findall(r"[a-zA-Z\']+", text)
                 if len(w) > 2 and w.lower() not in EN_STOP]
        for n in range(2, 5):
            for i in range(len(words) - n + 1):
                gram = " ".join(words[i:i+n])
                if gram not in existing_lower and gram not in ngrams:
                    ngrams[gram] = {"text": text, "label": label}
    if not ngrams:
        return
    try:
        scored = _xlmr_score(ngrams, HATE_ANCHORS_EN, threshold=0.72)
        for gram, sim in scored:
            if any(c["term"].lower() == gram.lower() for c in all_results):
                continue
            all_results.append({
                "term":              gram,
                "suggested_label":   ngrams[gram]["label"],
                "extraction_method": "english_xlmr",
                "score":             round(sim, 3),
                "source_text":       ngrams[gram]["text"][:300],
                "language":          "english",
            })
            print(f"  [english_xlmr] sim={sim:.3f} Kept: \'{gram}\'")
    except Exception as e:
        print(f"  [english_xlmr] Error: {e}")


def extract_candidates(hate_flagged_rows, normal_rows=None,
                       existing_terms=None, max_candidates=40):
    """LLM + XLMR extraction with language routing."""
    if not hate_flagged_rows:
        return []

    if existing_terms is None:
        try:
            from detector.models import LexiconEntry
            existing_terms = set(LexiconEntry.objects.values_list("term", flat=True))
        except Exception:
            existing_terms = set()
    existing_lower = {t.lower().strip() for t in existing_terms}

    by_lang = {"english": [], "amharic": [], "oromo": []}
    for r in hate_flagged_rows:
        by_lang[detect_language(r.get("text", ""))].append(r)

    all_results = []
    batch_size  = 20

    lang_batches = []
    for lang_key, lang_rows in [("english", by_lang["english"]),
                                 ("amharic", by_lang["amharic"]),
                                 ("oromo",   by_lang["oromo"])]:
        for start in range(0, len(lang_rows), batch_size):
            b = lang_rows[start:start + batch_size]
            if b:
                lang_batches.append((lang_key, b))

    for lang_key, batch in lang_batches:
        if lang_key != "english":
            continue

        # Check if Qwen is available
        _pn  = sys.modules.get("pattern_namer")
        _reg = getattr(_pn, "_model_registry", {}) if _pn else {}
        _qwen = "Qwen/Qwen2.5-1.5B-Instruct" in _reg

        if not _qwen:
            print(f"  [llm_extract] Qwen not cached -- using XLMR for English")
            _extract_english_terms_xlmr(batch, existing_lower, all_results)
            continue

        try:
            import torch
            tokenizer, model = _get_llm("english")
            print(f"  [llm_extract] Analyzing {len(batch)} texts (english)...")

            texts_block = chr(10).join(
                f"{i+1}. {r.get(chr(116)+chr(101)+chr(120)+chr(116), chr(39)+chr(39))[:250]}"
                for i, r in enumerate(batch)
            )
            texts_block = chr(10).join(
                f"{i+1}. {r.get('text', '')[:250]}" for i, r in enumerate(batch)
            )

            prompt = (
                "Extract the key hate-bearing PHRASES (2-5 words max) from these texts." + chr(10)
                + "Extract the specific slur, insult, or violent phrase -- NOT the whole sentence." + chr(10)
                + "Copy words exactly as written. Output a flat JSON array only." + chr(10) + chr(10)
                + "Example:" + chr(10)
                + "Input: These muslim savages must be expelled from our country" + chr(10)
                + "Output: [" + chr(34) + "muslim savages" + chr(34)
                + ", " + chr(34) + "must be expelled" + chr(34) + "]" + chr(10) + chr(10)
                + "Texts:" + chr(10) + texts_block + chr(10) + chr(10)
                + "Output (2-5 word phrases only, flat JSON array):"
            )

            messages = [{"role": "user", "content": prompt}]
            try:
                text_input = tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True)
            except Exception:
                text_input = prompt

            inputs = tokenizer(text_input, return_tensors="pt",
                               truncation=True, max_length=3500).to(model.device)
            with torch.no_grad():
                gen_ids = model.generate(**inputs, max_new_tokens=400, do_sample=False,
                                         pad_token_id=tokenizer.pad_token_id)
            raw = tokenizer.decode(gen_ids[0][inputs["input_ids"].shape[1]:],
                                   skip_special_tokens=True).strip()
            print(f"  [llm_extract] Response: {raw[:100]}")

            match = re.search(r"\[.*?\]", raw, re.DOTALL)
            if not match:
                continue
            try:
                parsed = json.loads(match.group(0))
            except json.JSONDecodeError:
                parsed = re.findall(r'"\'([^\']{3,})\"\'', raw)
            terms = []
            for item in parsed:
                if isinstance(item, list):
                    terms.extend([x for x in item if isinstance(x, str)])
                elif isinstance(item, str):
                    terms.append(item)

            batch_text_lower = " ".join(r.get("text", "").lower() for r in batch)
            for term in terms:
                if not isinstance(term, str) or len(term.strip()) < 3:
                    continue
                t_clean = term.strip()
                if t_clean.lower() in existing_lower:
                    continue
                if t_clean.lower() not in batch_text_lower:
                    print(f"  [llm_extract] Rejected hallucination: {t_clean}")
                    continue
                if len(t_clean.split()) == 1 and t_clean.lower() in _COMMON_SINGLES:
                    print(f"  [llm_extract] Rejected common single word: {t_clean}")
                    continue
                if len(t_clean.split()) > 6:
                    print(f"  [llm_extract] Rejected full sentence: {t_clean[:50]}")
                    continue
                if any(c["term"].lower() == t_clean.lower() for c in all_results):
                    continue
                source, label = "", "Identity-Based Hate"
                for r in batch:
                    if t_clean.lower() in r.get("text", "").lower():
                        source = r.get("text", "")[:300]
                        label  = r.get("label") or r.get("predicted_label", "Identity-Based Hate")
                        break
                all_results.append({
                    "term":              t_clean,
                    "suggested_label":   label,
                    "extraction_method": "llm_direct",
                    "score":             8.0,
                    "source_text":       source,
                    "language":          detect_language(t_clean),
                })
        except Exception as e:
            print(f"  [llm_extract] Error: {e}")
            _extract_english_terms_xlmr(batch, existing_lower, all_results)

    # XLMR extraction for Amharic and Oromo
    for lang_key, batch in lang_batches:
        if lang_key == "amharic":
            _extract_amharic_terms(batch, existing_lower, all_results)
        elif lang_key == "oromo":
            _extract_oromo_terms(batch, existing_lower, all_results)

    print(f"  [llm_extract] Total extracted: {len(all_results)}")
    return all_results[:max_candidates]