"""Build a semantic-vector pack (.vectors.zip) for a .slovarypack dictionary.

Reference implementation of the Slovary Pack Format v1, section 3
(https://slovary.vip/spec/). The output lets the dictionary participate in
the app's on-device meaning-based search.

Vectors are only comparable when passages and queries are encoded by the SAME
model. The app encodes queries with the E5-small INT8 ONNX runtime it
downloads as its "AI module"; this script encodes passages with those same
model files. Download them once:

    https://slovary.vip/packs/ai/model.onnx
    https://slovary.vip/packs/ai/tokenizer.json

Usage:
    python make_vectors.py --pack org.example.mydict.slovarypack \
        --model-dir ./ai_model --out org.example.mydict.vectors.zip

Dependencies (unlike make_pack.py this cannot be stdlib-only):
    pip install numpy onnxruntime tokenizers
    # optional progress bar: pip install tqdm

Passage construction (must match the app's corpus contract exactly):
    "passage: {TERM} | {TERM} | {plain-text definition snippet[:800]}"[:2000]
mean-pooled over the attention mask, L2-normalised, then INT8-quantised
per dimension:  code = floor((x - vmin) / vdiff * 255), clipped to 0..255;
the app decodes  x = vmin + (code + 0.5) * vdiff / 255.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import re
import sqlite3
import sys
import time
import zipfile
from pathlib import Path

import numpy as np

try:
    from tqdm import tqdm
except ImportError:  # progress bar is optional for the reference build
    def tqdm(it, **kw):
        return it

EMBEDDING_MODEL = "multilingual-e5-small-int8"
MAX_TOKENS = 512
SNIPPET_CHARS = 800
PASSAGE_CHARS = 2000

_TAG_RE = re.compile(r"<[^>]+>")
_ENT_RE = re.compile(r"&\w+;|&#\d+;")


def passage_text(term: str, definition: str) -> str:
    plain = _TAG_RE.sub(" ", definition)
    plain = _ENT_RE.sub(" ", plain)
    plain = " ".join(plain.split())[:SNIPPET_CHARS]
    return f"passage: {term} | {term} | {plain}"[:PASSAGE_CHARS]


def load_entries(pack_path: Path):
    db = sqlite3.connect("file:%s?mode=ro" % pack_path.as_posix(), uri=True)
    meta = dict(db.execute("SELECT key, value FROM pack_meta"))
    rows = db.execute(
        "SELECT local_id, term_upper, definition FROM entries "
        "ORDER BY local_id").fetchall()
    db.close()
    count = int(meta["entries_count"])
    if len(rows) != count:
        raise SystemExit(f"entries_count={count} but entries has {len(rows)} rows")
    if rows and (rows[0][0] != 1 or rows[-1][0] != count):
        raise SystemExit("local_id is not a dense 1..entries_count sequence")
    return meta, rows


class Encoder:
    """E5 passage/query encoder over the app's ONNX model files."""

    def __init__(self, model_dir: Path, batch_size: int):
        import onnxruntime as ort
        from tokenizers import Tokenizer

        self.batch = batch_size
        self.tok = Tokenizer.from_file(str(model_dir / "tokenizer.json"))
        self.tok.enable_truncation(max_length=MAX_TOKENS)
        # Padding is done by hand below, per batch: enabling it on the
        # tokenizer would pad every sequence to the longest of the WHOLE call,
        # i.e. 512 tokens for a dictionary-sized corpus.
        self.pad_id = self.tok.token_to_id("<pad>") or 1
        self.sess = ort.InferenceSession(
            str(model_dir / "model.onnx"), sess_options=ort.SessionOptions(),
            providers=ort.get_available_providers())
        self.input_names = {i.name for i in self.sess.get_inputs()}

    def encode(self, texts: list[str], progress: bool = False) -> np.ndarray:
        """Returns L2-normalised float32 embeddings, one row per text.

        Sequences are tokenised once, then grouped into batches of similar
        length: padding waste dominates the runtime on dictionary articles,
        whose lengths differ by an order of magnitude. Results are written
        back to the caller's original order.
        """
        encodings = self.tok.encode_batch(texts)
        token_ids = [e.ids for e in encodings]
        order = sorted(range(len(texts)), key=lambda i: len(token_ids[i]))
        out = None
        rng = range(0, len(order), self.batch)
        if progress:
            rng = tqdm(rng, desc="encode", unit="batch")
        for start in rng:
            idx = order[start:start + self.batch]
            width = max(len(token_ids[i]) for i in idx)
            ids = np.full((len(idx), width), self.pad_id, dtype=np.int64)
            mask = np.zeros((len(idx), width), dtype=np.int64)
            for r, i in enumerate(idx):
                seq = token_ids[i]
                ids[r, :len(seq)] = seq
                mask[r, :len(seq)] = 1
            feeds = {"input_ids": ids, "attention_mask": mask}
            if "token_type_ids" in self.input_names:
                feeds["token_type_ids"] = np.zeros_like(ids)
            hidden = self.sess.run(None, feeds)[0]  # (b, n, d)
            m = mask[:, :, None].astype(np.float32)
            vec = (hidden * m).sum(axis=1) / np.maximum(m.sum(axis=1), 1e-9)
            vec /= np.maximum(np.linalg.norm(vec, axis=1, keepdims=True), 1e-12)
            if out is None:
                out = np.empty((len(texts), vec.shape[1]), dtype=np.float32)
            out[idx] = vec.astype(np.float32)
        return out


def quantize(vectors: np.ndarray):
    vmin = vectors.min(axis=0)
    vdiff = np.maximum(vectors.max(axis=0) - vmin, 1e-9)
    codes = np.floor((vectors - vmin) / vdiff * 255.0)
    return np.clip(codes, 0, 255).astype(np.uint8), vmin, vdiff


def dequantize(codes: np.ndarray, vmin: np.ndarray, vdiff: np.ndarray):
    return vmin + (codes.astype(np.float32) + 0.5) * vdiff / 255.0


def self_check(enc: Encoder, rows, codes, vmin, vdiff, sample: int = 20) -> float:
    """Self-retrieval sanity: encoding an entry's own headword as a query must
    put an entry with that headword into the top-10. Returns the hit rate."""
    rng = np.random.default_rng(7)
    picks = rng.choice(len(rows), size=min(sample, len(rows)), replace=False)
    queries = [f"query: {rows[i][1]}" for i in picks]
    q = enc.encode(queries)
    decoded = dequantize(codes, vmin, vdiff)  # (count, d) float32
    scores = q @ decoded.T
    top = np.argsort(-scores, axis=1)[:, :10]
    hits = 0
    for qi, i in enumerate(picks):
        want = rows[i][1]
        if any(rows[j][1] == want for j in top[qi]):
            hits += 1
    return hits / len(picks)


def main() -> None:
    ap = argparse.ArgumentParser(
        description=__doc__.splitlines()[0],
        formatter_class=argparse.RawDescriptionHelpFormatter, epilog=__doc__)
    ap.add_argument("--pack", required=True, help="input .slovarypack")
    ap.add_argument("--model-dir", required=True,
                    help="directory with model.onnx + tokenizer.json")
    ap.add_argument("--out", required=True, help="output .vectors.zip")
    ap.add_argument("--batch", type=int, default=64)
    ap.add_argument("--skip-check", action="store_true",
                    help="skip the self-retrieval sanity check")
    args = ap.parse_args()

    t0 = time.time()
    pack_path = Path(args.pack)
    model_dir = Path(args.model_dir)
    for p in (pack_path, model_dir / "model.onnx", model_dir / "tokenizer.json"):
        if not p.exists():
            raise SystemExit(f"Missing input: {p}")

    meta, rows = load_entries(pack_path)
    print(f"{meta['pack_id']}: {len(rows)} entries")

    enc = Encoder(model_dir, args.batch)
    texts = [passage_text(term, definition) for _, term, definition in rows]
    vectors = enc.encode(texts, progress=True)
    codes, vmin, vdiff = quantize(vectors)

    if not args.skip_check:
        rate = self_check(enc, rows, codes, vmin, vdiff)
        print(f"self-retrieval top-10 hit rate: {rate:.0%}")
        if rate < 0.7:
            raise SystemExit("CHECK FAIL: self-retrieval below 70% - "
                             "wrong model files or a broken build")

    model_rev = hashlib.sha256(
        (model_dir / "model.onnx").read_bytes()).hexdigest()[:16]
    vec_meta = {
        "format_version": 1,
        "pack_id": meta["pack_id"],
        "entries_count": len(rows),
        "count": len(rows),
        "d": int(vectors.shape[1]),
        "embedding_model": EMBEDDING_MODEL,
        "model_revision": model_rev,
        "quantization": "int8-per-dim-minmax",
        "query_prefix": "query: ",
        "passage_prefix": "passage: ",
        "vmin": [float(x) for x in vmin],
        "vdiff": [float(x) for x in vdiff],
    }
    ids = np.arange(1, len(rows) + 1, dtype="<i8")

    out = Path(args.out)
    with zipfile.ZipFile(out, "w", zipfile.ZIP_DEFLATED) as z:
        z.writestr("vectors_meta.json",
                   json.dumps(vec_meta, ensure_ascii=False, sort_keys=True))
        z.writestr("vectors.sq8", codes.tobytes())
        z.writestr("ids.bin", ids.tobytes())
    mb = out.stat().st_size / 1048576
    print(f"Done in {time.time() - t0:.0f}s: {out} ({mb:.1f} MB)")


if __name__ == "__main__":
    sys.exit(main())