"""
embedder.py — Shared embedding model. Uses fastembed (ONNX, ~80MB RAM)
with fallback to sentence-transformers (PyTorch, ~500MB RAM).
"""

import os
from typing import Any

import numpy as np
from tqdm import tqdm

MODEL_NAME = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
# fastembed's files (~240 MB) live with the app, not in the temp dir it uses by
# default: that is empty on every Space start and after a reboot. The server
# build fills it (`python embedder.py`); if it is missing, fastembed downloads
# into it as before.
MODEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")

# Any: a TextEmbedding or a SentenceTransformer, both imported lazily; the
# two have different APIs, so callers dispatch on _backend.
_model: Any = None
_backend = None


def _init():
    global _model, _backend
    if _model is not None:
        return

    # Try fastembed first (lightweight ONNX)
    try:
        from fastembed import TextEmbedding
        _model = TextEmbedding(MODEL_NAME, cache_dir=MODEL_DIR)
        _backend = "fastembed"
        return
    except Exception:  # noqa: BLE001, S110 - any fastembed failure (import, download, ONNX) falls back below
        pass

    # Fall back to sentence-transformers (heavy PyTorch)
    try:
        from sentence_transformers import SentenceTransformer
        _model = SentenceTransformer(MODEL_NAME)
        _backend = "sentence-transformers"
        return
    except Exception:  # noqa: BLE001, S110 - reported by the RuntimeError below
        pass

    raise RuntimeError("Nelze načíst embedding model. Nainstaluj fastembed nebo sentence-transformers.")


def get_backend():
    _init()
    return _backend


def encode(text):
    """Encode a single text string. Returns numpy array."""
    _init()
    if _backend == "fastembed":
        return np.array(next(iter(_model.embed([text]))))
    else:
        return _model.encode(text)


def encode_batch(texts, show_progress=True):
    """Encode a list of texts. Returns numpy array of shape (n, dim).

    For large batches, prefers sentence-transformers (faster on GPU/CPU).
    Falls back to fastembed if sentence-transformers unavailable.
    """
    # For batch encoding, prefer sentence-transformers (much faster)
    try:
        from sentence_transformers import SentenceTransformer
        st_model = SentenceTransformer(MODEL_NAME)
        return st_model.encode(texts, show_progress_bar=show_progress, batch_size=256)
    except ImportError:
        pass

    # Fallback to fastembed (slower for batch but works without PyTorch)
    _init()
    if _backend == "fastembed":
        embeddings = list(
            tqdm(
                _model.embed(texts, batch_size=256),
                total=len(texts),
                desc="Generuji embeddingy",
                disable=not show_progress,
            )
        )
        return np.array(embeddings)
    else:
        return _model.encode(texts, show_progress_bar=show_progress, batch_size=256)


class Embedder:
    """Wrapper class compatible with retrieve.py interface (has .encode method)."""
    def encode(self, text):
        return encode(text)


if __name__ == "__main__":
    # Server builds run this to fetch the model into MODEL_DIR once.
    _init()
    print(f"embedder: {_backend} model ready in {MODEL_DIR}")