"""Local embedding with ONNX Runtime (low memory) or sentence-transformers fallback.

Default path uses ONNX Runtime + HuggingFace tokenizers for ~200MB RAM.
Falls back to sentence-transformers (~1.4GB RAM) if onnxruntime is not installed.

The model is loaded lazily on first use and auto-unloaded after idle timeout.
"""
import gc
import logging
import threading
import time
from pathlib import Path
from typing import Optional

import numpy as np

from desksearch.config import DEFAULT_EMBEDDING_MODEL

logger = logging.getLogger(__name__)

_DEFAULT_IDLE_TIMEOUT = 300  # 5 minutes

# HuggingFace ONNX model repo for all-MiniLM-L6-v2
_ONNX_MODEL_REPO = "sentence-transformers/all-MiniLM-L6-v2"
_ONNX_MODEL_FILE = "onnx/model.onnx"
_ONNX_DIMENSION = 384  # Known dimension for all-MiniLM-L6-v2


def _onnx_available() -> bool:
    """Check if onnxruntime is installed."""
    try:
        import onnxruntime  # noqa: F401
        return True
    except ImportError:
        return False


def _sentence_transformers_available() -> bool:
    """Check if sentence-transformers is installed."""
    try:
        import sentence_transformers  # noqa: F401
        return True
    except ImportError:
        return False


class Embedder:
    """Wraps an embedding model for generating embeddings.

    Prefers ONNX Runtime (low memory) and falls back to sentence-transformers.
    The model is loaded lazily, cached, and auto-unloaded after idle.
    """

    def __init__(
        self,
        model_name: str = DEFAULT_EMBEDDING_MODEL,
        idle_timeout: float = _DEFAULT_IDLE_TIMEOUT,
    ) -> None:
        self.model_name = model_name
        self.idle_timeout = idle_timeout

        self._model = None          # SentenceTransformer or None
        self._onnx_session = None   # onnxruntime.InferenceSession or None
        self._tokenizer = None      # AutoTokenizer or None
        self._backend: Optional[str] = None  # "onnx" or "sentence_transformers"
        self._lock = threading.Lock()
        self._last_used: float = 0.0
        self._unload_timer: Optional[threading.Timer] = None
        self._dimension: Optional[int] = None

    # ------------------------------------------------------------------
    # Model lifecycle
    # ------------------------------------------------------------------

    @property
    def is_loaded(self) -> bool:
        """Whether the embedding model is currently in memory."""
        return self._model is not None or self._onnx_session is not None

    @property
    def backend(self) -> Optional[str]:
        """Return which backend is active: 'onnx' or 'sentence_transformers'."""
        return self._backend

    def _get_onnx_model_path(self) -> Path:
        """Download/cache the ONNX model and return its path."""
        from huggingface_hub import hf_hub_download

        model_path = hf_hub_download(
            repo_id=_ONNX_MODEL_REPO,
            filename=_ONNX_MODEL_FILE,
        )
        return Path(model_path)

    def _ensure_loaded(self):
        """Load the model if not already loaded (thread-safe)."""
        if self._onnx_session is not None or self._model is not None:
            self._touch()
            return

        with self._lock:
            if self._onnx_session is not None or self._model is not None:
                self._touch()
                return

            t0 = time.perf_counter()

            if _onnx_available() and self.model_name == DEFAULT_EMBEDDING_MODEL:
                self._load_onnx()
            elif _sentence_transformers_available():
                self._load_sentence_transformers()
            else:
                raise RuntimeError(
                    "No embedding backend available. Install onnxruntime "
                    "(recommended, low memory) or sentence-transformers."
                )

            elapsed = (time.perf_counter() - t0) * 1000
            logger.info(
                "Model loaded in %.0fms via %s backend. Dimension: %d",
                elapsed, self._backend, self._dimension,
            )
            self._touch()

    def _load_onnx(self) -> None:
        """Load the ONNX model + lightweight tokenizer."""
        import onnxruntime as ort
        from tokenizers import Tokenizer

        logger.info("Loading ONNX embedding model: %s", self.model_name)

        # Download/cache the ONNX model
        model_path = self._get_onnx_model_path()

        # Configure ONNX Runtime for low memory
        sess_options = ort.SessionOptions()
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        sess_options.intra_op_num_threads = 2
        sess_options.inter_op_num_threads = 1

        self._onnx_session = ort.InferenceSession(
            str(model_path),
            sess_options=sess_options,
            providers=["CPUExecutionProvider"],
        )

        # Use lightweight tokenizers library (not transformers)
        self._tokenizer = Tokenizer.from_pretrained(_ONNX_MODEL_REPO)
        self._tokenizer.enable_padding(pad_id=0, pad_token="[PAD]")
        self._tokenizer.enable_truncation(max_length=256)
        self._dimension = _ONNX_DIMENSION
        self._backend = "onnx"

    def _load_sentence_transformers(self) -> None:
        """Load the sentence-transformers model (heavier, fallback)."""
        from sentence_transformers import SentenceTransformer

        logger.info("Loading sentence-transformers model: %s", self.model_name)
        self._model = SentenceTransformer(self.model_name)
        self._dimension = self._model.get_sentence_embedding_dimension()
        self._backend = "sentence_transformers"

    def _touch(self) -> None:
        """Record model usage and reset the auto-unload timer."""
        self._last_used = time.monotonic()
        self._schedule_unload()

    def _schedule_unload(self) -> None:
        """Schedule auto-unload after idle_timeout seconds."""
        if self._unload_timer is not None:
            self._unload_timer.cancel()
        if self.idle_timeout > 0:
            self._unload_timer = threading.Timer(self.idle_timeout, self._maybe_unload)
            self._unload_timer.daemon = True
            self._unload_timer.start()

    def _maybe_unload(self) -> None:
        """Unload the model if it has been idle long enough."""
        if not self.is_loaded:
            return
        elapsed = time.monotonic() - self._last_used
        if elapsed >= self.idle_timeout:
            logger.info(
                "Model idle for %.0fs (threshold %.0fs), unloading to free memory",
                elapsed, self.idle_timeout,
            )
            self.cooldown()

    def warmup(self) -> None:
        """Eagerly load the model and run a dummy embedding to warm up."""
        t0 = time.perf_counter()
        _ = self.embed(["warmup"])
        elapsed = (time.perf_counter() - t0) * 1000
        logger.info("Embedder warmup complete in %.0fms", elapsed)

    def cooldown(self) -> None:
        """Explicitly unload the model to free memory."""
        with self._lock:
            if not self.is_loaded:
                return

            self._model = None
            self._onnx_session = None
            self._tokenizer = None
            gc.collect()

            try:
                import torch
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            except ImportError:
                pass

            logger.info("Embedding model unloaded (%s backend), memory freed", self._backend)

    # ------------------------------------------------------------------
    # Embedding
    # ------------------------------------------------------------------

    @property
    def dimension(self) -> int:
        """Return the embedding dimension of the model."""
        if self._dimension is not None:
            return self._dimension
        self._ensure_loaded()
        return self._dimension

    def _mean_pool_and_normalize(
        self, token_embeddings: np.ndarray, attention_mask: np.ndarray
    ) -> np.ndarray:
        """Mean pooling over token embeddings, then L2-normalize."""
        # token_embeddings: (batch, seq_len, hidden)
        # attention_mask: (batch, seq_len)
        mask_expanded = attention_mask[:, :, np.newaxis].astype(np.float32)
        summed = np.sum(token_embeddings * mask_expanded, axis=1)
        counts = np.clip(mask_expanded.sum(axis=1), a_min=1e-9, a_max=None)
        pooled = summed / counts
        # L2 normalize
        norms = np.linalg.norm(pooled, axis=1, keepdims=True)
        norms = np.clip(norms, a_min=1e-9, a_max=None)
        return (pooled / norms).astype(np.float32)

    def _embed_onnx(self, texts: list[str], batch_size: int) -> np.ndarray:
        """Embed using ONNX Runtime with lightweight tokenizers."""
        all_embeddings = []

        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            encoded = self._tokenizer.encode_batch(batch)

            input_ids = np.array([e.ids for e in encoded], dtype=np.int64)
            attention_mask = np.array([e.attention_mask for e in encoded], dtype=np.int64)
            token_type_ids = np.zeros_like(input_ids, dtype=np.int64)

            feeds = {
                "input_ids": input_ids,
                "attention_mask": attention_mask,
                "token_type_ids": token_type_ids,
            }

            outputs = self._onnx_session.run(None, feeds)
            # outputs[0] is token_embeddings (batch, seq_len, hidden)
            token_embeddings = outputs[0]
            pooled = self._mean_pool_and_normalize(
                token_embeddings, attention_mask.astype(np.float32)
            )
            all_embeddings.append(pooled)

        return np.vstack(all_embeddings)

    def _embed_sentence_transformers(self, texts: list[str], batch_size: int) -> np.ndarray:
        """Embed using sentence-transformers."""
        embeddings = self._model.encode(
            texts,
            batch_size=batch_size,
            show_progress_bar=False,
            convert_to_numpy=True,
            normalize_embeddings=True,
        )
        return embeddings.astype(np.float32)

    def embed(self, texts: list[str], batch_size: int = 64) -> np.ndarray:
        """Embed a list of text strings.

        Args:
            texts: List of text strings to embed.
            batch_size: Number of texts to embed at once.

        Returns:
            numpy array of shape (len(texts), dimension).
        """
        if not texts:
            return np.array([], dtype=np.float32).reshape(0, self.dimension)

        self._ensure_loaded()

        if self._backend == "onnx":
            return self._embed_onnx(texts, batch_size)
        else:
            return self._embed_sentence_transformers(texts, batch_size)

    def embed_query(self, query: str) -> np.ndarray:
        """Embed a single query string.

        Returns:
            1D numpy array of shape (dimension,).
        """
        return self.embed([query])[0]
