"""Embedding pipeline -- all-MiniLM-L6-v2 via ONNX Runtime.

No PyTorch dependency.  ~96MB model file.  <15ms per embedding on CPU.
Downloads the ONNX model from Hugging Face on first use.

Requires ``onnxruntime``, ``numpy``, and ``tokenizers`` (all in ``[smart]`` tier).
"""

from __future__ import annotations

import hashlib
import logging
import shutil
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Guarded imports -- core tier code must not crash when deps are absent
# ---------------------------------------------------------------------------

EMBEDDINGS_AVAILABLE: bool = False

try:
    import numpy as np
    import onnxruntime as ort

    EMBEDDINGS_AVAILABLE = True
except ImportError:
    np = None  # type: ignore[assignment]
    ort = None  # type: ignore[assignment]

# ``tokenizers`` is optional but strongly preferred for proper tokenization.
_TOKENIZERS_AVAILABLE: bool = False
try:
    from tokenizers import Tokenizer as _HFTokenizer  # type: ignore[import-untyped]

    _TOKENIZERS_AVAILABLE = True
except ImportError:
    _HFTokenizer = None  # type: ignore[assignment,misc]


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_HF_REPO = "sentence-transformers/all-MiniLM-L6-v2"
_ONNX_FILENAME = "model.onnx"
_TOKENIZER_FILENAME = "tokenizer.json"

# URLs for direct download from Hugging Face (no ``huggingface_hub`` needed)
_ONNX_URL = f"https://huggingface.co/{_HF_REPO}/resolve/main/onnx/{_ONNX_FILENAME}"
_TOKENIZER_URL = f"https://huggingface.co/{_HF_REPO}/resolve/main/{_TOKENIZER_FILENAME}"

# Expected SHA-256 of the ONNX model so we can verify integrity.  Set to
# ``None`` to skip verification (useful when the upstream model updates).
_ONNX_SHA256: str | None = None


# ============================================================================
# SimpleTokenizer
# ============================================================================


class SimpleTokenizer:
    """Lightweight tokenizer for the embedding pipeline.

    Uses the ``tokenizers`` library (Rust-backed, fast) when the vocab file
    is available.  Falls back to a basic word-splitting heuristic that
    produces *reasonable* embeddings for routing -- but not for NLP tasks.
    """

    # Special token IDs for the BERT-style vocabulary
    _PAD_ID = 0
    _CLS_ID = 101
    _SEP_ID = 102
    _UNK_ID = 100

    def __init__(self, vocab_path: Path | None = None) -> None:
        self._tokenizer: _HFTokenizer | None = None
        self._using_fast: bool = False

        if vocab_path is not None and vocab_path.is_file() and _TOKENIZERS_AVAILABLE:
            try:
                self._tokenizer = _HFTokenizer.from_file(str(vocab_path))
                self._using_fast = True
                logger.debug("Loaded fast tokenizer from %s", vocab_path)
            except Exception:
                logger.warning("Failed to load tokenizer from %s -- falling back to simple tokenizer", vocab_path)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def encode_batch(self, texts: list[str], max_length: int = 256) -> dict[str, Any]:
        """Encode *texts* into ``input_ids``, ``attention_mask``, and ``token_type_ids``.

        Returns numpy arrays shaped ``(N, max_length)`` ready for the ONNX model.
        """
        if not EMBEDDINGS_AVAILABLE:
            raise RuntimeError("numpy is required for tokenization -- install llmhost[smart]")

        if self._using_fast and self._tokenizer is not None:
            return self._encode_fast(texts, max_length)
        return self._encode_simple(texts, max_length)

    @property
    def is_fast(self) -> bool:
        """``True`` when using the Rust-backed ``tokenizers`` library."""
        return self._using_fast

    # ------------------------------------------------------------------
    # Fast path (tokenizers library)
    # ------------------------------------------------------------------

    def _encode_fast(self, texts: list[str], max_length: int) -> dict[str, Any]:
        """Encode using the ``tokenizers`` library for proper WordPiece output."""
        assert self._tokenizer is not None

        self._tokenizer.enable_truncation(max_length=max_length)
        self._tokenizer.enable_padding(length=max_length, pad_id=self._PAD_ID, pad_token="[PAD]")

        encoded = self._tokenizer.encode_batch(texts)

        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)

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "token_type_ids": token_type_ids,
        }

    # ------------------------------------------------------------------
    # Simple fallback (no external tokenizer)
    # ------------------------------------------------------------------

    def _encode_simple(self, texts: list[str], max_length: int) -> dict[str, Any]:
        """Best-effort encoding using whitespace splitting + character hashing.

        This will *not* produce the same token IDs as the real WordPiece
        vocabulary, but cosine similarity between embeddings produced this
        way is still informative enough for kNN routing decisions.
        """
        batch_ids: list[list[int]] = []
        batch_mask: list[list[int]] = []

        for text in texts:
            # Simple split and hash each token to a vocabulary-range ID
            words = text.lower().split()
            token_ids: list[int] = [self._CLS_ID]
            for word in words:
                if len(token_ids) >= max_length - 1:  # leave room for [SEP]
                    break
                # Hash word to a stable pseudo-token ID in [999, 30521)
                hashed = int(hashlib.md5(word.encode(), usedforsecurity=False).hexdigest(), 16) % 29522 + 999
                token_ids.append(hashed)
            token_ids.append(self._SEP_ID)

            # Pad / truncate
            attention = [1] * len(token_ids)
            pad_len = max_length - len(token_ids)
            if pad_len > 0:
                token_ids.extend([self._PAD_ID] * pad_len)
                attention.extend([0] * pad_len)
            else:
                token_ids = token_ids[:max_length]
                attention = attention[:max_length]

            batch_ids.append(token_ids)
            batch_mask.append(attention)

        input_ids = np.array(batch_ids, dtype=np.int64)
        attention_mask = np.array(batch_mask, dtype=np.int64)
        token_type_ids = np.zeros_like(input_ids, dtype=np.int64)

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "token_type_ids": token_type_ids,
        }


# ============================================================================
# EmbeddingPipeline
# ============================================================================


class EmbeddingPipeline:
    """Generate embeddings using all-MiniLM-L6-v2 via ONNX Runtime.

    No PyTorch needed.  ~96MB model file.  <15ms per embedding.
    Downloads model on first use to ``~/.llmhosts/models/minilm/``.

    Usage::

        pipeline = EmbeddingPipeline()
        await pipeline.initialize()          # downloads model if missing
        vec = pipeline.embed("Hello world")  # (384,) float32
    """

    MODEL_NAME: str = "all-MiniLM-L6-v2"
    EMBEDDING_DIM: int = 384
    MAX_SEQ_LENGTH: int = 256

    def __init__(self, model_dir: Path | None = None) -> None:
        if not EMBEDDINGS_AVAILABLE:
            raise RuntimeError(
                "EmbeddingPipeline requires numpy and onnxruntime. Install them with: pip install llmhosts[smart]"
            )
        self._model_dir: Path = model_dir or Path.home() / ".llmhosts" / "models" / "minilm"
        self._session: ort.InferenceSession | None = None
        self._tokenizer: SimpleTokenizer | None = None

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    async def initialize(self) -> None:
        """Load or download the ONNX model and initialise the inference session.

        Sequence:
        1. Ensure ``model_dir`` exists.
        2. If ``model.onnx`` is missing, attempt to download from Hugging Face.
        3. Load the ONNX session with CPU execution provider.
        4. Initialise the tokenizer (prefer ``tokenizer.json``, fall back to simple).
        """
        self._model_dir.mkdir(parents=True, exist_ok=True)

        onnx_path = self._model_dir / _ONNX_FILENAME
        tokenizer_path = self._model_dir / _TOKENIZER_FILENAME

        # Download if not present
        if not onnx_path.is_file():
            await self._download_model(onnx_path, tokenizer_path)

        if not onnx_path.is_file():
            raise FileNotFoundError(
                f"ONNX model not found at {onnx_path}. "
                f"Download it manually:\n"
                f"  mkdir -p {self._model_dir}\n"
                f"  curl -L -o {onnx_path} {_ONNX_URL}\n"
                f"  curl -L -o {tokenizer_path} {_TOKENIZER_URL}"
            )

        # Verify integrity if hash is known
        if _ONNX_SHA256 is not None:
            self._verify_checksum(onnx_path, _ONNX_SHA256)

        # Create ONNX inference session (CPU only)
        sess_options = ort.SessionOptions()
        sess_options.inter_op_num_threads = 1
        sess_options.intra_op_num_threads = 2
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

        self._session = ort.InferenceSession(
            str(onnx_path),
            sess_options=sess_options,
            providers=["CPUExecutionProvider"],
        )
        logger.info("Loaded ONNX model from %s", onnx_path)

        # Initialise tokenizer
        self._tokenizer = SimpleTokenizer(vocab_path=tokenizer_path)
        if self._tokenizer.is_fast:
            logger.info("Using fast (Rust) tokenizer")
        else:
            logger.warning(
                "Using simple fallback tokenizer -- install 'tokenizers' and download "
                "tokenizer.json for best quality: curl -L -o %s %s",
                tokenizer_path,
                _TOKENIZER_URL,
            )

    # ------------------------------------------------------------------
    # Embedding API
    # ------------------------------------------------------------------

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

        Returns
        -------
        numpy.ndarray
            Shape ``(384,)`` float32 L2-normalised embedding vector.
        """
        result = self.embed_batch([text])
        return np.asarray(result[0])

    def embed_batch(self, texts: list[str]) -> np.ndarray:
        """Embed multiple texts in one forward pass.

        Returns
        -------
        numpy.ndarray
            Shape ``(N, 384)`` float32 L2-normalised embedding matrix.
        """
        if self._session is None or self._tokenizer is None:
            raise RuntimeError("Pipeline not initialised -- call await pipeline.initialize() first")

        if not texts:
            return np.empty((0, self.EMBEDDING_DIM), dtype=np.float32)

        # Tokenize
        encoded = self._tokenizer.encode_batch(texts, max_length=self.MAX_SEQ_LENGTH)

        # Run ONNX inference
        outputs = self._session.run(
            None,
            {
                "input_ids": encoded["input_ids"],
                "attention_mask": encoded["attention_mask"],
                "token_type_ids": encoded["token_type_ids"],
            },
        )

        # outputs[0] is the last hidden state: (batch, seq_len, hidden_dim)
        token_embeddings: np.ndarray = outputs[0]

        # Mean pooling with attention mask
        pooled = self._mean_pool(token_embeddings, encoded["attention_mask"])

        # L2 normalise for cosine similarity via inner product
        normalised = self._normalize(pooled)

        return normalised

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _mean_pool(self, embeddings: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
        """Mean pooling over token embeddings, respecting the attention mask.

        Parameters
        ----------
        embeddings:
            Token-level embeddings, shape ``(batch, seq_len, dim)``.
        attention_mask:
            Binary mask, shape ``(batch, seq_len)``.

        Returns
        -------
        numpy.ndarray
            Sentence embeddings, shape ``(batch, dim)``.
        """
        # Expand mask to match embedding dimensions: (batch, seq_len, 1)
        mask_expanded = np.expand_dims(attention_mask, axis=-1).astype(np.float32)

        # Zero out padding token embeddings and sum
        sum_embeddings = np.sum(embeddings * mask_expanded, axis=1)

        # Clamp denominator to avoid division by zero
        sum_mask = np.clip(mask_expanded.sum(axis=1), a_min=1e-9, a_max=None)

        return np.asarray(sum_embeddings / sum_mask)

    @staticmethod
    def _normalize(embeddings: np.ndarray) -> np.ndarray:
        """L2-normalise each embedding vector.

        After normalisation, ``np.dot(a, b)`` equals cosine similarity.
        """
        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
        norms = np.clip(norms, a_min=1e-12, a_max=None)
        return np.asarray(embeddings / norms)

    # ------------------------------------------------------------------
    # Model download
    # ------------------------------------------------------------------

    async def _download_model(self, onnx_path: Path, tokenizer_path: Path) -> None:
        """Download the ONNX model and tokenizer from Hugging Face.

        Uses ``urllib`` so there are no extra dependencies beyond the stdlib.
        Falls back to a helpful error message if the download fails (e.g. no
        internet access).
        """
        import urllib.error
        import urllib.request

        logger.info("Downloading %s ONNX model to %s ...", self.MODEL_NAME, self._model_dir)

        for url, dest in [(_ONNX_URL, onnx_path), (_TOKENIZER_URL, tokenizer_path)]:
            if dest.is_file():
                continue
            tmp = dest.with_suffix(".tmp")
            try:
                logger.info("  Fetching %s ...", url)
                urllib.request.urlretrieve(url, str(tmp))  # nosec B310
                shutil.move(str(tmp), str(dest))
                logger.info("  Saved %s (%s bytes)", dest.name, dest.stat().st_size)
            except (urllib.error.URLError, OSError) as exc:
                tmp.unlink(missing_ok=True)
                logger.warning("Download failed for %s: %s", url, exc)
                logger.info(
                    "You can download manually:\n  curl -L -o %s %s",
                    dest,
                    url,
                )

    @staticmethod
    def _verify_checksum(path: Path, expected_sha256: str) -> None:
        """Verify the SHA-256 checksum of a downloaded file."""
        sha = hashlib.sha256()
        with path.open("rb") as fh:
            while True:
                chunk = fh.read(1 << 20)  # 1 MiB
                if not chunk:
                    break
                sha.update(chunk)
        actual = sha.hexdigest()
        if actual != expected_sha256:
            raise ValueError(
                f"Checksum mismatch for {path.name}: expected {expected_sha256[:16]}..., got {actual[:16]}..."
            )

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def is_ready(self) -> bool:
        """``True`` when the ONNX session is loaded and the tokenizer is available."""
        return self._session is not None and self._tokenizer is not None

    @property
    def model_dir(self) -> Path:
        """Directory containing the ONNX model files."""
        return self._model_dir
