"""Standalone ONNX inference for Bashkir person-name recognition.

The graph returns logits at SentencePiece positions. Training assigned one BIO
label to the first piece of each original word, so this runtime does the same.
"""
from __future__ import annotations

import json
import re
from pathlib import Path

import numpy as np
import onnxruntime as ort
import sentencepiece as spm

TAGS = ("O", "B-PER", "I-PER")
WORD_RE = re.compile(r"[A-Za-zА-Яа-яЁёӘәҒғҘҙҠҡҢңӨөҪҫҮүҺһҚқІі][\w'\-]*", re.UNICODE)


class BashkirPersonNER:
    def __init__(self, model_dir: str | Path = ".", precision: str = "int8",
                 providers: list[str] | None = None):
        self.directory = Path(model_dir)
        config = json.loads((self.directory / "config.json").read_text(encoding="utf-8"))
        if precision not in config["models"]:
            raise ValueError(f"precision must be one of {sorted(config['models'])}")
        self.max_length = int(config["max_position_embeddings"])
        self.bos_id = int(config["bos_token_id"])
        self.eos_id = int(config["eos_token_id"])
        self.sp = spm.SentencePieceProcessor(
            model_file=str(self.directory / "spm_bashkir_bert_16k.model")
        )
        model = self.directory / config["models"][precision]
        self.session = ort.InferenceSession(
            str(model), providers=providers or ["CPUExecutionProvider"]
        )

    def predict_tokens(self, tokens: list[str]) -> tuple[list[str], bool]:
        """Return word-level BIO tags and whether input exceeded the context limit.

        The caller supplies the same token boundaries used for gold evaluation.
        On truncation, the returned tag list covers only its leading words.
        """
        if not tokens:
            return [], False
        ids = [self.bos_id]
        positions = []
        for word in tokens:
            pieces = list(self.sp.encode(word, out_type=int)) or [self.sp.unk_id()]
            if len(ids) + len(pieces) + 1 > self.max_length:
                break
            positions.append(len(ids))
            ids.extend(pieces)
        if not positions:
            return [], True
        ids.append(self.eos_id)
        logits = self.session.run(
            ["logits"], {"input_ids": np.asarray([ids], dtype=np.int64)}
        )[0][0]
        return [TAGS[int(logits[pos].argmax())] for pos in positions], len(positions) < len(tokens)

    def predict(self, text: str) -> dict:
        """Return person spans with character offsets into the original text."""
        if not isinstance(text, str):
            raise TypeError("text must be a string")
        matches = list(WORD_RE.finditer(text))
        tags, truncated = self.predict_tokens([m.group() for m in matches])
        entities = []
        start = None
        end = None
        for match, tag in zip(matches, tags):
            if tag == "B-PER":
                if start is not None:
                    entities.append({"text": text[start:end], "start": start, "end": end})
                start, end = match.span()
            elif tag == "I-PER":
                # Current benchmark convention: an orphan I-PER starts a span.
                if start is None:
                    start = match.start()
                end = match.end()
            elif start is not None:
                entities.append({"text": text[start:end], "start": start, "end": end})
                start = end = None
        if start is not None:
            entities.append({"text": text[start:end], "start": start, "end": end})
        return {"entities": entities, "truncated": truncated}


if __name__ == "__main__":
    import sys

    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8")
    ner = BashkirPersonNER(Path(__file__).resolve().parent)
    for sample in sys.argv[1:]:
        print(json.dumps(ner.predict(sample), ensure_ascii=False))