"""Text-Embedding-Daemon für die Inhaltssuche (multilingual-e5-small, ONNX/CPU).

Warum kein sentence-transformers: das zöge torch + transformers ins venv, ~2 GB
auf einer Karte, auf der schon whisper, sherpa und ein 5-GB-LLM liegen. Hier
reichen `tokenizers` (Rust, schon installiert) und `onnxruntime` — beides ist
bereits für whisper/NPU da, es kommt keine neue Abhängigkeit dazu.

Protokoll wie spk_embed: eine JSON-Zeile rein, eine JSON-Zeile raus.
    {"texts": ["query: ...", "passage: ..."]}  ->  {"vecs": [[...], [...]]}
Fehler:                                           {"error": "..."}

Aufruf: python embed_text.py /opt/aime-models/e5-small
"""
import json
import sys
import os

import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer

MAX_LEN = 512


def log(msg):
    print(msg, file=sys.stderr, flush=True)


def main():
    if len(sys.argv) < 2:
        log("usage: embed_text.py <model-dir>")
        return 1
    d = sys.argv[1]
    tok = Tokenizer.from_file(os.path.join(d, "tokenizer.json"))
    tok.enable_truncation(max_length=MAX_LEN)
    tok.enable_padding()

    so = ort.SessionOptions()
    # Vier A55-Kerne, passend zum taskset in emb.js. Mehr Threads als Kerne
    # kosten hier nur Kontextwechsel.
    so.intra_op_num_threads = 4
    so.inter_op_num_threads = 1
    sess = ort.InferenceSession(os.path.join(d, "model.onnx"), so, providers=["CPUExecutionProvider"])
    want = {i.name for i in sess.get_inputs()}

    print("READY", flush=True)

    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            texts = json.loads(line)["texts"]
            encs = tok.encode_batch(texts)
            ids = np.array([e.ids for e in encs], dtype=np.int64)
            mask = np.array([e.attention_mask for e in encs], dtype=np.int64)
            feed = {"input_ids": ids, "attention_mask": mask}
            # Manche ONNX-Exporte verlangen token_type_ids, andere kennen es nicht.
            if "token_type_ids" in want:
                feed["token_type_ids"] = np.zeros_like(ids)
            feed = {k: v for k, v in feed.items() if k in want}
            out = sess.run(None, feed)[0]          # [B, T, H] last_hidden_state
            # Mean-Pooling über die echten Token (Padding darf nicht mitzählen),
            # danach L2 — so ist das Skalarprodukt direkt der Kosinus.
            m = mask[..., None].astype(np.float32)
            vec = (out * m).sum(axis=1) / np.clip(m.sum(axis=1), 1e-9, None)
            vec = vec / np.clip(np.linalg.norm(vec, axis=1, keepdims=True), 1e-9, None)
            print(json.dumps({"vecs": [[round(float(x), 6) for x in v] for v in vec]}), flush=True)
        except Exception as e:                      # noqa: BLE001 - Daemon darf nie sterben
            print(json.dumps({"error": f"{type(e).__name__}: {e}"}), flush=True)
    return 0


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