backend/src/services/rerankService.js

import { config } from "../config/index.js";
import { chatCompletion } from "./ollamaClient.js";

// módulo-folha (sem importar de chatChain.js/atendimentosQueryService.js), pra poder ser
// reusado tanto no rerank de documentos quanto no de atendimentos sem dependência circular
// (mesmo motivo de hydeService.js).
const RERANK_SYSTEM_PROMPT = [
  "Você reordena trechos de documentos internos por relevância à pergunta do usuário.",
  'Responda APENAS com um objeto JSON no formato {"order": [...]}, contendo TODOS os índices',
  '(0-based) de cada trecho, em ordem decrescente de relevância. Exemplo: {"order": [2,0,3,1]}.',
  "Não inclua explicações, texto extra ou markdown."
].join("\n");

const RERANK_SNIPPET_LENGTH = 300;

function buildRerankPrompt(query, hits) {
  const listing = hits
    .map((h, i) => `[${i}] ${h.text.slice(0, RERANK_SNIPPET_LENGTH)}`)
    .join("\n\n");
  return `Pergunta: ${query}\n\nTrechos:\n${listing}`;
}

function parseRerankOrder(content, len) {
  try {
    const cleaned = String(content ?? "")
      .trim()
      .replace(/^```json\s*/i, "")
      .replace(/^```\s*/, "")
      .replace(/```$/, "")
      .trim();
    const parsed = JSON.parse(cleaned);
    const arr = Array.isArray(parsed) ? parsed : Array.isArray(parsed?.order) ? parsed.order : null;
    if (!arr) return null;

    const seen = new Set();
    const valid = [];
    for (const v of arr) {
      const idx = Number(v);
      if (Number.isInteger(idx) && idx >= 0 && idx < len && !seen.has(idx)) {
        seen.add(idx);
        valid.push(idx);
      }
    }
    if (valid.length === 0) return null;
    for (let i = 0; i < len; i++) if (!seen.has(i)) valid.push(i);
    return valid;
  } catch {
    return null;
  }
}

export async function rerankHits(query, hits, { topN, numCtx = config.llm.numCtx } = {}) {
  if (hits.length <= 1) return hits;
  const limit = topN ?? hits.length;
  try {
    const { content } = await chatCompletion({
      messages: [
        { role: "system", content: RERANK_SYSTEM_PROMPT },
        { role: "user", content: buildRerankPrompt(query, hits) }
      ],
      format: "json",
      options: { temperature: 0.1, num_ctx: numCtx, num_predict: config.llm.auxNumPredict }
    });
    const order = parseRerankOrder(content, hits.length);
    if (!order) console.warn("[rerankService] rerankHits: resposta não parseável, mantendo ordem original");
    const reordered = order ? order.map((i) => hits[i]) : hits;
    return reordered.slice(0, limit);
  } catch (err) {
    console.warn("[rerankService] rerankHits falhou, mantendo ordem original:", err.message);
    return hits.slice(0, limit);
  }
}