"""Arc Tagger: text-only streaming tagger (about 10 M parameters, CPU) for ASR word streams.

    from arctag import ArcTagger
    tg = ArcTagger.from_pretrained()                 # WhissleAI/arc-tagger from the Hub, or a local dir
    out = tg.tag("Hi, this is Dana from Acme. Can you move my appointment to Friday?")

    st = tg.stream(domain="talk")                    # streaming: words in as the ASR emits them
    st.push_words(["hi", "this", "is"], final=False) # interim
    st.push_words(["hi,", "this", "is", "Dana."], revise_from=0, final=True)
    st.endpoint()
    res = st.finalize()

CPU only. The model code is the streaming engine in arctag/_core.
"""
import importlib.util, json, os, sys

__version__ = "0.4.1"
_HERE = os.path.dirname(os.path.abspath(__file__))
_CORE = os.path.join(_HERE, "_core")
_EN = None


def _engine(home):
    """import the vendored engine once (module-level path constants resolve against `home`)"""
    global _EN
    if _EN is None:
        os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
        os.environ.setdefault("ARCTAG_HOME", home)
        os.environ["ARCTAG_ROLE"] = "offline"
        if _CORE not in sys.path:
            sys.path.insert(0, _CORE)
        spec = importlib.util.spec_from_file_location("arctag_engine", os.path.join(_CORE, "engine.py"))
        _EN = importlib.util.module_from_spec(spec); sys.modules["arctag_engine"] = _EN
        spec.loader.exec_module(_EN)
    return _EN


def _driver():
    from . import driver
    return driver


class ArcTagger:
    def __init__(self, home, EN, loaded, cfg):
        self.home, self.EN, self.loaded, self.config = home, EN, loaded, cfg

    @classmethod
    def from_pretrained(cls, path="WhissleAI/arc-tagger", revision=None, thresholds="refit"):
        """path: release dir (model.safetensors, config.json, tok/, ontology/, thresholds/) or a Hugging Face repo id.
        thresholds: "refit" (default: thresholds/display.json, plus thresholds/events.json and span_head.json when present)
        or "none" (the default thresholds in config.json only; engine display floors still apply)."""
        if not os.path.isdir(path):
            from huggingface_hub import snapshot_download
            path = snapshot_download(path, revision=revision, allow_patterns=["*.json", "*.safetensors", "*.npy", "*.py", "*.txt", "*.md",
                                                                               "LICENSE*", "tok/*", "ontology/*", "thresholds/*", "arctag/*"])
        home = os.path.abspath(path)
        EN = _engine(home)
        AT = sys.modules["arc_text"]; TT = sys.modules["tagtable"]
        AT.BPE8K = os.path.join(home, "tok", "bpe8k", "tokenizer.json")
        cfg = json.load(open(os.path.join(home, "config.json")))
        ext = os.path.join(home, cfg.get("onto_ext") or "ontology/ext_v7_aff.json")
        T = TT.TextTable(r8=os.path.join(home, "ontology", "tags_v1_r8.json"),
                         text_tags=os.path.join(home, "ontology", "text_tags.json"), ext=ext)
        T.load_E(os.path.join(home, "ontology", "head_E_r8_bge_small.npy"))
        if T.sha != cfg.get("ontology_sha"):
            raise RuntimeError(f"ontology sha {T.sha} != config {cfg.get('ontology_sha')}")
        loaded = _load_model(EN, home, cfg, T, thresholds)
        return cls(home, EN, loaded, cfg)

    def stream(self, domain="auto", tier="T0"):
        """a streaming session (driver.Stream): push_words(words, final=True, revise_from=None, endpoint=False), endpoint(),
        finalize() -> result dict"""
        return _driver().Stream(self.EN, self.loaded, domain=domain, tier=tier, meta=self._meta())

    def tag(self, x, domain="auto", tier="T0"):
        """x: text (str), a list of words (str) or word dicts {"w", "t0", "t1", "conf", "spk", "final"}. Utterances end after
        words ending in . ? ! Returns the structured result (see README "Output")."""
        return _driver().tag(self.EN, self.loaded, x, domain=domain, tier=tier, meta=self._meta())

    def _meta(self):
        return {"model": "arc-tagger", "version": __version__}


def _read_thr(display, events):
    """thresholds/display.json {"thresholds": {tag@readout: {domain|_all: thr}}} and thresholds/events.json {"tier", "hide",
    "thresholds"}"""
    pref, pstep, ev, evstep = {}, None, {}, None
    if os.path.exists(display):
        pref, pstep = json.load(open(display)).get("thresholds") or {}, "display"
    if os.path.exists(events):
        ev, evstep = json.load(open(events)), "events"
    return pref, pstep, ev, evstep


def _load_model(EN, home, cfg, T, thresholds="refit"):
    """load model.safetensors and the threshold tables of a release dir"""
    from safetensors.torch import load_file
    AT = sys.modules["arc_text"]; SA = sys.modules["stream_arc"]
    step = None
    m = AT.ArcText(T, scratch=cfg.get("scratch") or AT.SCRATCH_10M)
    sd = load_file(os.path.join(home, "model.safetensors"))
    SA.add_modules_for(m, sd)
    missing, unexpected = m.load_state_dict(sd, strict=False)
    if unexpected:
        raise RuntimeError(f"weights without a module: {unexpected[:6]}")
    bad = [k for k in missing if not k.startswith(("head.E",))]
    if bad:
        import warnings
        warnings.warn(f"{len(bad)} missing keys (e.g. {bad[:3]})")
    m.eval()
    own = cfg.get("thresholds") or {}
    pref, pstep, ev, evstep = {}, None, {}, None
    if thresholds == "refit":
        d = os.path.join(home, "thresholds")
        pref, pstep, ev, evstep = _read_thr(os.path.join(d, "display.json"), os.path.join(d, "events.json"))
        sh = os.path.join(d, "span_head.json")          # span-head thresholds per tier
        if os.path.exists(sh):
            ts = {t: dict(v) for t, v in (cfg.get("thresholds_span") or {}).items()}
            for t, v in (json.load(open(sh)).get("thresholds_span") or {}).items():
                ts.setdefault(t, {}).update(v)
            cfg = dict(cfg, thresholds_span=ts)
    elif thresholds != "none":
        raise ValueError("thresholds must be 'refit' or 'none'")
    thr = EN.Thresholds(own, {}, None, pref=pref, pref_step=pstep)
    thr.set_events(ev, evstep)
    L = EN.Loaded(step, home, m, T, thr, cfg)
    fams = cfg.get("span_head_families")              # span families decoded from the span head (others: typed span decode)
    if fams is not None and L.head_tags:
        L.head_use, L.head_why = set(fams) & set(L.head_tags), "config"
    return L