from __future__ import annotations
import json
import os
import numpy as np
import torch
from safetensors.torch import load_file
from data import real_pii
from data.taxonomy import N_PII_TAGS, TAG_TO_ID, TAXONOMY
from modeling import (
    FELAModerationV2,
    ModerationConfig,
    encode_text,
    load_model,
    pad_batch,
)

OPENAI11 = TAXONOMY[:11]


def heldout_taxonomy(n=600):
    from datasets import load_dataset

    ds = load_dataset("google/civil_comments", split="validation", streaming=True)
    thr = 0.5
    m = {
        "hate": "identity_attack",
        "harassment": "insult",
        "harassment_threatening": "threat",
        "sexual": "sexual_explicit",
        "violence": "threat",
    }
    out = []
    for ex in ds:
        if not ex.get("text"):
            continue
        lab = [0.0] * 11
        msk = [0.0] * 11
        for i, cat in enumerate(OPENAI11):
            col = m.get(cat)
            if col is not None and ex.get(col) is not None:
                lab[i] = 1.0 if ex[col] >= thr else 0.0
                msk[i] = 1.0
        out.append((ex["text"], lab, msk))
        if len(out) >= n:
            break
    return out


def heldout_pii(n=200):
    rows = list(real_pii.load_nemotron(max_rows=n))
    return rows


JIGSAW_COLS = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]


def heldout_jigsaw(n=3000):
    from data.toxicity import _iter_jigsaw

    out = []
    for text, ex in _iter_jigsaw(streaming=True, max_rows=n, want_eval=True):
        out.append((text, [float(int(ex[c])) for c in JIGSAW_COLS], [1.0] * 6))
    return out or None


def eval_jigsaw(model, rows, bs=32):
    from sklearn.metrics import roc_auc_score

    texts = [r[0] for r in rows]
    y = np.array([r[1] for r in rows])
    model.eval()
    probs = []
    with torch.no_grad():
        for i in range(0, len(texts), bs):
            ids, mask = _batch(texts[i : i + bs])
            probs.append(torch.sigmoid(model(ids, mask, task="jigsaw")).cpu().numpy())
    p = np.concatenate(probs, 0)
    res = {}
    for j, c in enumerate(JIGSAW_COLS):
        yy = y[:, j]
        if yy.min() == yy.max():
            continue
        res[c] = round(float(roc_auc_score(yy, p[:, j])), 4)
    if res:
        res["mean"] = round(sum(res.values()) / len(res), 4)
    return res


def _batch(texts, max_len=512):
    seqs = [encode_text(t, max_len)[0] for t in texts]
    return pad_batch(seqs, max_len)


def taxonomy_probs(model, texts, bs=32):
    model.eval()
    probs = []
    with torch.no_grad():
        for i in range(0, len(texts), bs):
            ids, mask = _batch(texts[i : i + bs])
            logits = model(ids, mask, task="taxonomy")[:, :11]
            probs.append(torch.sigmoid(logits).cpu().numpy())
    return np.concatenate(probs, 0)


def eval_taxonomy(model, rows):
    from sklearn.metrics import roc_auc_score

    texts = [r[0] for r in rows]
    labels = np.array([r[1] for r in rows])
    masks = np.array([r[2] for r in rows])
    p = taxonomy_probs(model, texts)
    res = {}
    for i, cat in enumerate(OPENAI11):
        sel = masks[:, i] > 0
        y = labels[sel, i]
        if sel.sum() < 20 or y.min() == y.max():
            continue
        res[cat] = round(float(roc_auc_score(y, p[sel, i])), 4)
    return res


def eval_pii(model, rows, max_len=512):
    model.eval()
    correct = tot = 0
    with torch.no_grad():
        from data.bio import spans_to_bio

        for text, spans in rows:
            ids, tags = spans_to_bio(text, spans, max_len, TAG_TO_ID)
            idt = torch.tensor([ids])
            mask = torch.ones_like(idt)
            pred = model(idt, mask, task="pii")[0].argmax(-1).tolist()
            for t, pr in zip(tags, pred):
                tot += 1
                correct += int(t == pr)
    return round(correct / max(tot, 1), 4)


def load_int8(model, tier_path, scales_path):
    q = load_file(tier_path)
    scales = json.load(open(scales_path))
    state = {}
    for k, v in q.items():
        if v.dtype == torch.int8 and k in scales:
            state[k] = v.float() * scales[k]
        else:
            state[k] = v
    model.load_state_dict(state, strict=False)
    return model


def heldout_clf(task, n=400, seed=99):
    import importlib

    mod = importlib.import_module(f"data.{task}")
    rows = []
    synth = getattr(mod, f"synth_{task}", None)
    if synth is not None:
        rows.extend(list(synth(n // 2, seed=seed)))
    real = getattr(mod, f"load_{task}_examples", None)
    if real is not None:
        try:
            rows.extend(list(real(max_rows=n // 2)))
        except Exception:
            pass
    return rows


def eval_clf_head(model, task, rows, label_names):
    from sklearn.metrics import roc_auc_score

    if not rows:
        return {}
    texts = [r["text"] for r in rows]
    y = np.array([r["labels"] for r in rows])
    msk = np.array([r["mask"] for r in rows])
    model.eval()
    probs = []
    with torch.no_grad():
        for i in range(0, len(texts), 32):
            ids, mask = _batch(texts[i : i + 32])
            probs.append(torch.sigmoid(model(ids, mask, task=task)).cpu().numpy())
    p = np.concatenate(probs, 0)
    res = {}
    for j, name in enumerate(label_names):
        sel = msk[:, j] > 0
        yy = y[sel, j]
        if sel.sum() < 20 or yy.min() == yy.max():
            continue
        res[name] = round(float(roc_auc_score(yy, p[sel, j])), 4)
    return res


def head_scorecard(model):
    from data.taxonomy import V2_HEADS

    card = {}
    for task, labels in V2_HEADS.items():
        rows = heldout_clf(task)
        aur = eval_clf_head(model, task, rows, labels)
        mean = round(sum(aur.values()) / len(aur), 4) if aur else None
        gate = (
            "SHIP"
            if mean and mean >= 0.8
            else "HOLD"
            if mean and mean >= 0.7
            else "DEFER"
        )
        card[task] = {
            "per_label_auroc": aur,
            "mean_auroc": mean,
            "gate": gate,
            "n_eval": len(rows),
        }
    return card


_TOP10 = {
    "en": "en_US",
    "zh": "zh_CN",
    "hi": "hi_IN",
    "es": "es_ES",
    "fr": "fr_FR",
    "ar": "ar_AA",
    "bn": "bn_BD",
    "pt": "pt_BR",
    "ru": "ru_RU",
    "ja": "ja_JP",
}


def _assemble_vals(segments):
    parts, spans, blen = ([], [], 0)
    for seg in segments:
        if isinstance(seg, tuple):
            ent, val = seg
            vb = str(val).encode("utf-8")
            spans.append((blen, blen + len(vb), ent))
            parts.append(str(val))
            blen += len(vb)
        else:
            parts.append(seg)
            blen += len(seg.encode("utf-8"))
    return ("".join(parts), spans)


def heldout_pii_lang(locale, n=60):
    from faker import Faker

    try:
        f = Faker(locale)
    except Exception:
        return []
    rows = []
    for _ in range(n):
        try:
            segs = [
                "",
                ("PERSON", f.name()),
                " | ",
                ("EMAIL", f.email()),
                " | ",
                ("PHONE", f.phone_number()),
                " | ",
                ("STREET_ADDRESS", f.street_address()),
            ]
        except Exception:
            continue
        rows.append(_assemble_vals(segs))
    return rows


def multilingual_pii(model):
    out = {}
    for lang, loc in _TOP10.items():
        rows = heldout_pii_lang(loc)
        out[lang] = {
            "pii_token_acc": eval_pii(model, rows) if rows else None,
            "n": len(rows),
        }
    return out


def build(n_tax):
    return FELAModerationV2(ModerationConfig(n_pii_tags=N_PII_TAGS), n_tax=n_tax)


def main():
    import argparse

    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--artifact", required=True, help="dir with model.safetensors + tier_full_int8"
    )
    ap.add_argument(
        "--parent", default="lowdown-labs/fela-moderator", help="parent v1 repo/path"
    )
    ap.add_argument("--n-tax", type=int, default=len(TAXONOMY))
    args = ap.parse_args()
    tax_rows = heldout_taxonomy()
    pii_rows = heldout_pii()
    jig_rows = heldout_jigsaw()
    print(
        f"held-out: taxonomy={len(tax_rows)} pii={len(pii_rows)} jigsaw={('unreachable' if jig_rows is None else len(jig_rows))}"
    )
    report = {}
    new = build(args.n_tax)
    new.load_state_dict(load_file(os.path.join(args.artifact, "model.safetensors")))
    report["new_fp32"] = {
        "taxonomy_auroc": eval_taxonomy(new, tax_rows),
        "pii_token_acc": eval_pii(new, pii_rows),
    }
    if jig_rows:
        report["new_fp32"]["jigsaw_auroc"] = eval_jigsaw(new, jig_rows)
    report["v2_head_scorecard"] = head_scorecard(new)
    report["multilingual_pii"] = multilingual_pii(new)
    try:
        int8 = build(args.n_tax)
        load_int8(
            int8,
            os.path.join(args.artifact, "tier_full_int8.safetensors"),
            os.path.join(args.artifact, "tier_full_scales.json"),
        )
        report["new_int8"] = {
            "taxonomy_auroc": eval_taxonomy(int8, tax_rows),
            "pii_token_acc": eval_pii(int8, pii_rows),
        }
    except Exception as e:
        report["new_int8"] = {"error": repr(e)}
    try:
        parent = load_model(args.parent, strict=False)
        report["parent_v1"] = {
            "taxonomy_auroc": eval_taxonomy(parent, tax_rows),
            "pii_token_acc": eval_pii(parent, pii_rows),
        }
        if jig_rows:
            report["parent_v1"]["jigsaw_auroc"] = eval_jigsaw(parent, jig_rows)
    except Exception as e:
        report["parent_v1"] = {"error": repr(e)}
    print("PARITY_EVAL", json.dumps(report, indent=2))
    return report


if __name__ == "__main__":
    main()