from __future__ import annotations
import functools
import concurrent.futures
import threading
import os

try:
    import numpy as np
except Exception:  # pragma: no cover
    np = None  # type: ignore

try:
    from sentence_transformers import SentenceTransformer
except Exception:  # model optional in CI/dev
    SentenceTransformer = None  # type: ignore

_pool = concurrent.futures.ThreadPoolExecutor(max_workers=2)
_model_lock = threading.Lock()
_model_obj = None


def _load_model():
    global _model_obj
    with _model_lock:
        if _model_obj is None and SentenceTransformer is not None:
            # Use a small, widely available model
            _model_obj = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
    return _model_obj


def _norm(s: str) -> str:
    return (s or "").strip().lower()


@functools.lru_cache(maxsize=4096)
def _cached_score(a_norm: str, b_norm: str) -> float:
    m = _load_model()
    if m is None or np is None:
        return 0.0
    # Encode in the model thread (outside main loop)
    embs = m.encode([a_norm, b_norm])
    v1, v2 = np.asarray(embs[0]), np.asarray(embs[1])
    denom = (np.linalg.norm(v1) * np.linalg.norm(v2)) or 1.0
    return float(np.dot(v1, v2) / denom)


def score_pair(a: str, b: str, timeout_ms: int = None) -> float:
    """Return [0,1] cosine similarity for a,b strings with hard timeout.

    If the model is unavailable or times out, returns 0.0.
    """
    # Allow env override for timeout
    if timeout_ms is None:
        try:
            timeout_ms = int(os.getenv("SEMANTIC_TIMEOUT_MS", "150"))
        except Exception:
            timeout_ms = 150

    a_n, b_n = _norm(a), _norm(b)
    if not a_n or not b_n:
        return 0.0
    fut = _pool.submit(_cached_score, a_n, b_n)
    try:
        val = fut.result(timeout=max(1, timeout_ms) / 1000.0)
        # Clamp to [0,1]
        return max(0.0, min(1.0, float(val)))
    except Exception:
        return 0.0
