"""Embedding backends: ONNX Runtime (fast) > sentence-transformers > Ollama."""
from __future__ import annotations
import json, logging, math, urllib.request
from pathlib import Path
from ._constants import EMBED_MODEL, OLLAMA_ENDPOINT

log = logging.getLogger("pulse.rule_embeddings")
_ST_MODEL = None
_ONNX_SESSION = None
_ONNX_TOKENIZER = None
_ONNX_DIR = Path(__file__).parent / "onnx_model"


def _embed_onnx(texts: list[str]) -> list[list[float]]:
    """Fast ONNX Runtime embedding — loads in ~200ms vs 14s for PyTorch."""
    global _ONNX_SESSION, _ONNX_TOKENIZER
    import numpy as np
    if _ONNX_SESSION is None:
        onnx_path = _ONNX_DIR / "model.onnx"
        if not onnx_path.exists():
            raise FileNotFoundError(f"ONNX model not found at {onnx_path}")
        import onnxruntime as ort
        _ONNX_SESSION = ort.InferenceSession(
            str(onnx_path), providers=["CPUExecutionProvider"])
        from tokenizers import Tokenizer
        _ONNX_TOKENIZER = Tokenizer.from_file(str(_ONNX_DIR / "tokenizer.json"))
        _ONNX_TOKENIZER.enable_padding(length=128)
        _ONNX_TOKENIZER.enable_truncation(max_length=128)
    encoded = _ONNX_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.array([e.type_ids for e in encoded], dtype=np.int64)
    outputs = _ONNX_SESSION.run(None, {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "token_type_ids": token_type_ids,
    })
    # Mean pooling over token embeddings (masked)
    hidden = outputs[0]  # (batch, seq, 384)
    mask_expanded = attention_mask[:, :, np.newaxis].astype(np.float32)
    summed = (hidden * mask_expanded).sum(axis=1)
    counts = mask_expanded.sum(axis=1).clip(min=1e-9)
    pooled = summed / counts
    # L2 normalize
    norms = np.linalg.norm(pooled, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    normalized = pooled / norms
    return normalized.tolist()


def _embed_ollama(texts: list[str], model: str = EMBED_MODEL) -> list[list[float]]:
    url = f"{OLLAMA_ENDPOINT}/api/embed"
    payload = json.dumps({"model": model, "input": texts}).encode()
    req = urllib.request.Request(url, data=payload, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=120) as resp:
        data = json.loads(resp.read().decode())
    embeddings = data.get("embeddings", [])
    if len(embeddings) != len(texts):
        raise ValueError(f"Expected {len(texts)} embeddings, got {len(embeddings)}")
    return embeddings


def _embed_sentence_transformers(texts: list[str]) -> list[list[float]]:
    global _ST_MODEL
    if _ST_MODEL is None:
        import io, os, sys
        os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
        os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
        os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
        # Force local-only model load — skip HuggingFace HTTP checks.
        # The model is already cached locally; network HEAD requests add ~20s.
        os.environ.setdefault("HF_HUB_OFFLINE", "1")
        os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
        old_out, old_err = sys.stdout, sys.stderr
        sys.stdout = io.StringIO(); sys.stderr = io.StringIO()
        try:
            from sentence_transformers import SentenceTransformer
            _ST_MODEL = SentenceTransformer("all-MiniLM-L6-v2", local_files_only=True)
        finally:
            sys.stdout = old_out; sys.stderr = old_err
    return [v.tolist() for v in _ST_MODEL.encode(texts, show_progress_bar=False)]


def embed_texts(texts: list[str]) -> list[list[float]]:
    """Embed texts: ONNX Runtime (fast) > sentence-transformers > Ollama.

    ONNX loads in ~200ms and infers in ~5ms. sentence-transformers loads
    PyTorch (14s import chain). Ollama requires a running server.
    """
    if not texts:
        return []
    # Fast path: ONNX Runtime (~200ms cold, ~5ms warm)
    if _ONNX_DIR.exists() and (_ONNX_DIR / "model.onnx").exists():
        try:
            return _embed_onnx(texts)
        except Exception as e:
            log.debug("ONNX embedding failed (%s), trying alternatives", e)
    # Medium path: sentence-transformers (14s cold, ~30ms warm)
    try:
        return _embed_sentence_transformers(texts)
    except (ImportError, Exception) as e:
        log.debug("sentence-transformers failed (%s), trying Ollama", e)
    # Slow path: Ollama (requires running server)
    try:
        return _embed_ollama(texts)
    except Exception as e:
        raise RuntimeError(f"No embedding backend available: {e}")


def cosine_similarity(a: list[float], b: list[float]) -> float:
    dot = sum(x * y for x, y in zip(a, b))
    na = sum(x * x for x in a) ** 0.5
    nb = sum(x * x for x in b) ** 0.5
    return 0.0 if na == 0 or nb == 0 else dot / (na * nb)
