import { join } from "node:path";
import {
  AutoTokenizer,
  AutoModelForSequenceClassification,
  env
} from "@huggingface/transformers";
import { createLogger } from "../../logger.js";
import { getConfig } from "../../config.js";
const log = createLogger("reranker");
function wrapLoadError(err) {
  const message = err instanceof Error ? err.message : String(err);
  const looksLikeMissingModel = /tokenizer_class|tokenizer_config|preprocessor_config|fetch failed|ENOTFOUND|ECONNREFUSED|ETIMEDOUT|ENETUNREACH/i.test(
    message
  );
  if (looksLikeMissingModel) {
    return new Error(
      `Reranker model not downloaded \u2014 run \`wigolo warmup\` (cause: ${message})`
    );
  }
  return new Error(`Failed to load reranker model: ${message}`);
}
class TransformersRerankProvider {
  tokenizer = null;
  model = null;
  loadPromise = null;
  modelId;
  constructor() {
    this.modelId = "Xenova/ms-marco-MiniLM-L-6-v2";
  }
  async warmup() {
    await this.load();
  }
  load() {
    if (this.tokenizer && this.model) {
      return Promise.resolve({ tokenizer: this.tokenizer, model: this.model });
    }
    if (this.loadPromise) return this.loadPromise;
    log.info("Loading rerank model", { modelId: this.modelId });
    const cacheDir = join(getConfig().dataDir, "transformers");
    env.cacheDir = cacheDir;
    this.loadPromise = Promise.all([
      AutoTokenizer.from_pretrained(this.modelId),
      AutoModelForSequenceClassification.from_pretrained(this.modelId)
    ]).then(([tokenizer, model]) => {
      this.tokenizer = tokenizer;
      this.model = model;
      return { tokenizer, model };
    }).catch((err) => {
      this.loadPromise = null;
      throw wrapLoadError(err);
    });
    return this.loadPromise;
  }
  async rerank(query, candidates, topK = candidates.length) {
    if (candidates.length === 0) return [];
    const { tokenizer, model } = await this.load();
    const queries = candidates.map(() => query);
    const docs = candidates.map((c) => c.text);
    const inputs = tokenizer(queries, {
      text_pair: docs,
      padding: true,
      truncation: true
    });
    const outputs = await model(inputs);
    const logits = outputs.logits;
    const stride = logits.dims.length >= 2 ? logits.dims[1] : 1;
    const data = logits.data;
    const scored = candidates.map((c, i) => ({
      id: c.id,
      score: Number(data[i * stride])
    }));
    return scored.sort((a, b) => b.score - a.score).slice(0, topK);
  }
  // Release the underlying ONNX session before process exit. Without this,
  // the runtime's worker threads race during C++ destructor teardown and
  // surface as `mutex lock failed: Invalid argument` on macOS.
  async dispose() {
    const model = this.model;
    this.model = null;
    this.tokenizer = null;
    this.loadPromise = null;
    if (!model) return;
    try {
      const m = model;
      if (typeof m.dispose === "function") await m.dispose();
    } catch (err) {
      log.debug("reranker dispose failed", { error: err instanceof Error ? err.message : String(err) });
    }
  }
}
export {
  TransformersRerankProvider
};
//# sourceMappingURL=transformers-rerank-provider.js.map