#!/usr/bin/env python3

import json
import os
import signal
import sys
from pathlib import Path
from typing import Any

# Keep stdout clean for the PHP caller. Diagnostics go to stderr only.
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
os.environ.setdefault("HF_HUB_ETAG_TIMEOUT", "5")
os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "10")

MODEL_LOAD_TIMEOUT_SECONDS = int(os.environ.get("RETRIEX_EMBEDDING_MODEL_LOAD_TIMEOUT_SECONDS", "30"))


def empty() -> None:
    print("[]")
    sys.exit(0)


def debug(message: str) -> None:
    print(message, file=sys.stderr, flush=True)


# ---------------------------------------------------------
# Positional args (aligned with PHP client exec call)
# ---------------------------------------------------------
# 1 query
# 2 limit
# 3 index_path
# 4 meta_path
# 5 model
# ---------------------------------------------------------
if len(sys.argv) < 6:
    empty()

query = sys.argv[1]

try:
    limit = int(sys.argv[2])
except Exception:
    limit = 5

index_path = Path(sys.argv[3]).resolve()
meta_path = Path(sys.argv[4]).resolve()
model_name = sys.argv[5]

model_override = os.environ.get("RETRIEX_EMBEDDING_MODEL_PATH", "").strip()
if model_override:
    model_name = model_override

# ---------------------------------------------------------
# Dependency checks
# ---------------------------------------------------------
try:
    import faiss
except Exception:
    empty()

try:
    from sentence_transformers import SentenceTransformer
except Exception:
    empty()

# ---------------------------------------------------------
# File checks
# ---------------------------------------------------------
if limit <= 0:
    empty()

if not index_path.is_file() or not meta_path.is_file():
    empty()

# ---------------------------------------------------------
# Load model with timeout
# ---------------------------------------------------------
def load_model(model_name_or_path: str) -> SentenceTransformer:
    def timeout_handler(_signum: int, _frame: Any) -> None:
        raise TimeoutError("tag search embedding model load timed out")

    previous_handler = signal.signal(signal.SIGALRM, timeout_handler)
    signal.alarm(MODEL_LOAD_TIMEOUT_SECONDS)

    try:
        return SentenceTransformer(model_name_or_path)
    finally:
        signal.alarm(0)
        signal.signal(signal.SIGALRM, previous_handler)


try:
    model = load_model(model_name)
except Exception as exc:
    debug(f"Unable to load tag search embedding model '{model_name}': {exc}")
    empty()

# ---------------------------------------------------------
# Load index + meta
# ---------------------------------------------------------
try:
    index = faiss.read_index(str(index_path))
except Exception:
    empty()

try:
    with open(meta_path, "r", encoding="utf-8") as f:
        ids = json.load(f)
except Exception:
    empty()

if not isinstance(ids, list) or len(ids) == 0:
    empty()

# ---------------------------------------------------------
# Embed & search
# ---------------------------------------------------------
try:
    qvec = model.encode([query], normalize_embeddings=True, show_progress_bar=False)
    scores, idxs = index.search(qvec, limit)
except Exception:
    empty()

out = []
for score, idx in zip(scores[0], idxs[0]):
    if idx is None or idx < 0 or idx >= len(ids):
        continue
    out.append({
        "tag_id": str(ids[idx]),
        "score": float(score),
    })

print(json.dumps(out))
sys.exit(0)
