"""Portable Unicode char_wb preprocessing and sparse ONNX inference."""
import hashlib
import json
import re
from collections import Counter
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent


def sha256(path):
    with open(path, "rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def char_wb(text, ngram_range):
    # Match sklearn's _char_wb_ngrams, including short-word handling.
    text = re.sub(r"\s\s+", " ", text.lower())
    for word in text.split():
        word = " " + word + " "
        for n in range(ngram_range[0], ngram_range[1] + 1):
            offset = 0
            yield word[offset:offset + n]
            while offset + n < len(word):
                offset += 1
                yield word[offset:offset + n]
            if offset == 0:
                break


WORD_RE = re.compile(r"(?u)\b\w\w+\b")


def word_ngrams(text, ngram_range):
    tokens = WORD_RE.findall(text.lower())
    for n in range(ngram_range[0], ngram_range[1] + 1):
        for i in range(len(tokens) - n + 1):
            yield " ".join(tokens[i:i + n])


class LanguageIdentifier:
    def __init__(self, model_dir=None, batch_size=256):
        import onnxruntime as ort

        if model_dir is not None:
            self.model_dir = Path(model_dir)
        else:
            self.model_dir = ROOT / "model" if (ROOT / "model" / "META.json").is_file() else ROOT
        config_path = self.model_dir / "config.json"
        self.config = json.loads(config_path.read_text(encoding="utf-8")) if config_path.is_file() else None
        self.meta = json.loads((self.model_dir / "META.json").read_text(encoding="utf-8"))
        for name in ("model.onnx", "vectorizer.json"):
            if name in self.meta.get("sha256", {}):
                digest = self.meta["sha256"][name]
                if sha256(self.model_dir / name) != digest:
                    raise ValueError(f"LID artifact checksum mismatch: {name}")
        cfg = json.loads((self.model_dir / "vectorizer.json").read_text(encoding="utf-8"))
        self.char_vocab = cfg.get("char_vocabulary", cfg.get("vocabulary", {}))
        self.char_ngram_range = tuple(cfg.get("char_ngram_range", cfg.get("ngram_range", (2, 5))))
        self.word_vocab = cfg.get("word_vocabulary", {})
        self.word_ngram_range = tuple(cfg.get("word_ngram_range", (1, 2)))
        self.classes = np.asarray(self.meta["classes"])
        if self.config is not None:
            expected = {"classes": self.classes.tolist(), "model_file": "model.onnx",
                        "vectorizer_file": "vectorizer.json", "metadata_file": "META.json"}
            if any(self.config.get(k) != v for k, v in expected.items()):
                raise ValueError("LID config does not match model artifacts")
        if batch_size < 1:
            raise ValueError("batch_size must be positive")
        self.batch_size = batch_size
        options = ort.SessionOptions()
        options.intra_op_num_threads = 1
        self.session = ort.InferenceSession(str(self.model_dir / "model.onnx"), options,
                                            providers=["CPUExecutionProvider"])

    def features(self, texts):
        char_rows = []
        word_rows = []
        for text in texts:
            if not isinstance(text, str):
                raise TypeError("LID input must contain strings")
            c_counts = Counter(self.char_vocab[g] for g in char_wb(text, self.char_ngram_range)
                               if g in self.char_vocab)
            char_rows.append(c_counts)
            if self.word_vocab:
                w_counts = Counter(self.word_vocab[w] for w in word_ngrams(text, self.word_ngram_range)
                                   if w in self.word_vocab)
                word_rows.append(w_counts)

        w_c = max(1, max(map(len, char_rows), default=0))
        c_ids = np.zeros((len(char_rows), w_c), dtype=np.int64)
        c_counts = np.zeros((len(char_rows), w_c), dtype=np.float32)
        for i, row in enumerate(char_rows):
            if row:
                c_ids[i, :len(row)] = list(row.keys())
                c_counts[i, :len(row)] = list(row.values())

        if self.word_vocab:
            w_w = max(1, max(map(len, word_rows), default=0))
            w_ids = np.zeros((len(word_rows), w_w), dtype=np.int64)
            w_counts = np.zeros((len(word_rows), w_w), dtype=np.float32)
            for i, row in enumerate(word_rows):
                if row:
                    w_ids[i, :len(row)] = list(row.keys())
                    w_counts[i, :len(row)] = list(row.values())
            return {"char_ids": c_ids, "char_counts": c_counts, "word_ids": w_ids, "word_counts": w_counts}
        return {"char_ids": c_ids, "char_counts": c_counts}

    def predict_proba(self, texts):
        texts = list(texts)
        if not texts:
            return np.empty((0, len(self.classes)), dtype=np.float32)
        results = []
        for start in range(0, len(texts), self.batch_size):
            batch_texts = texts[start:start + self.batch_size]
            results.append(self.session.run(["probabilities"], self.features(batch_texts))[0])
        return np.concatenate(results) if results else np.empty((0, len(self.classes)), dtype=np.float32)

    def predict(self, texts):
        texts = list(texts)
        if not texts:
            return np.empty(0, dtype=str)
        return self.classes[self.predict_proba(texts).argmax(axis=1)]

    def describe(self):
        return {"name": self.meta["id"], "version": self.meta["version"],
                "sha256": self.meta["sha256"], "backend": "onnxruntime",
                "decision_rule": "argmax probabilities; classes ordered as in model passport"}