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

# Default model location relative to this file (src/ → project root → models/)
_DEFAULT_MODEL_PATH = Path(__file__).parent.parent / "models" / "Xenova" / "all-MiniLM-L6-v2"


class Embedder:
    def __init__(self, path=None):
        path = Path(path) if path is not None else _DEFAULT_MODEL_PATH
        if not (path / "tokenizer.json").exists() or not (path / "model.onnx").exists():
            import sys
            root_dir = Path(__file__).parent.parent
            if str(root_dir / "scripts") not in sys.path:
                sys.path.append(str(root_dir / "scripts"))
            from download import download
            download("Xenova/all-MiniLM-L6-v2", dest=root_dir / "models")

        self.tokenizer = Tokenizer.from_file(str(path / "tokenizer.json"))
        self.session = ort.InferenceSession(
            str(path / "model.onnx"), providers=["CPUExecutionProvider"]
        )
        self.input_names = {inp.name for inp in self.session.get_inputs()}

    def encode(self, text, normalize=True):
        return self.encode_batch([text], normalize=normalize)[0]

    def encode_batch(self, texts, normalize=True):
        self.tokenizer.enable_padding()
        encoded = self.tokenizer.encode_batch(texts)
        feed = {}
        if "input_ids" in self.input_names:
            feed["input_ids"] = np.array([e.ids for e in encoded], dtype=np.int64)
        if "attention_mask" in self.input_names:
            feed["attention_mask"] = np.array(
                [e.attention_mask for e in encoded], dtype=np.int64
            )
        if "token_type_ids" in self.input_names:
            feed["token_type_ids"] = np.array(
                [e.type_ids for e in encoded], dtype=np.int64
            )
        hidden = self.session.run(None, feed)[0]
        mask = feed["attention_mask"][..., None]
        pooled = (hidden * mask).sum(axis=1) / mask.sum(axis=1)
        if normalize:
            pooled = pooled / np.linalg.norm(pooled, axis=1, keepdims=True)
        return pooled