import { Injectable, Logger } from '@nestjs/common';

/**
 * Servicio de re-ranking: cross-encoder que ordena chunks por
 * relevancia REAL a la pregunta, no por similitud vectorial.
 *
 * Soporta:
 *  - cohere   (API, modelo recomendado: rerank-multilingual-v3.0)
 *  - jina     (API, modelo: jina-reranker-v2-base-multilingual)
 *  - local    (Ollama u otro endpoint compatible)
 *
 * Diseñado para ser agnóstico del proveedor — la lógica de RAG
 * en ayuda-ia.service.ts no debería saber cuál se está usando.
 */

export interface RerankConfig {
  proveedor: 'cohere' | 'jina' | 'local';
  modelo: string;                // ej: 'rerank-multilingual-v3.0'
  apiKey: string | null;
  baseUrl?: string | null;       // solo para 'local'
  topK: number;                  // chunks a devolver tras el rerank
  minScore: number;              // umbral mínimo (0-1) — chunks por debajo se descartan
}

export interface RerankInput<T> {
  query: string;
  documents: Array<{
    /** Texto del chunk que el reranker va a evaluar contra la query. */
    text: string;
    /** Payload original — se devuelve intacto en el resultado. */
    payload: T;
  }>;
}

export interface RerankResult<T> {
  /** Índice original (0..N-1) en el array recibido. */
  index: number;
  /** Score de relevancia del cross-encoder (0-1). */
  score: number;
  /** Payload original sin cambios. */
  payload: T;
}

@Injectable()
export class RerankerService {
  private readonly logger = new Logger(RerankerService.name);

  /**
   * Re-ordena documentos por relevancia. Si `config.proveedor` no está soportado
   * o falla la API, devuelve los documentos en el orden recibido con score 0
   * (degradación graceful — el caller decide qué hacer).
   */
  async rerank<T>(config: RerankConfig, input: RerankInput<T>): Promise<RerankResult<T>[]> {
    if (input.documents.length === 0) return [];

    try {
      if (config.proveedor === 'cohere') {
        return this.rerankCohere(config, input);
      }
      if (config.proveedor === 'jina') {
        return this.rerankJina(config, input);
      }
      if (config.proveedor === 'local') {
        return this.rerankLocal(config, input);
      }
      this.logger.warn(`Proveedor reranker desconocido: ${config.proveedor}. Devuelvo orden original.`);
      return this.fallback(input);
    } catch (err: any) {
      this.logger.error(`Rerank falló (${config.proveedor}): ${err?.message ?? err}. Devuelvo orden original.`);
      return this.fallback(input);
    }
  }

  // ───────────────────────────────────────────────────────────────
  // Cohere Rerank API
  //   POST https://api.cohere.com/v2/rerank
  //   body: { model, query, documents: string[], top_n }
  //   docs: https://docs.cohere.com/reference/rerank
  // ───────────────────────────────────────────────────────────────
  private async rerankCohere<T>(
    config: RerankConfig,
    input: RerankInput<T>,
  ): Promise<RerankResult<T>[]> {
    if (!config.apiKey) {
      throw new Error('Cohere requiere API key configurada.');
    }
    const url = (config.baseUrl ?? 'https://api.cohere.com').replace(/\/$/, '') + '/v2/rerank';
    const body = {
      model: config.modelo,
      query: input.query,
      documents: input.documents.map((d) => d.text),
      top_n: Math.min(config.topK, input.documents.length),
    };
    const resp = await fetch(url, {
      method: 'POST',
      headers: {
        'Authorization': `Bearer ${config.apiKey}`,
        'Content-Type': 'application/json',
        'Accept': 'application/json',
      },
      body: JSON.stringify(body),
    });
    if (!resp.ok) {
      const text = await resp.text().catch(() => '');
      throw new Error(`Cohere ${resp.status}: ${text.slice(0, 200)}`);
    }
    const data = (await resp.json()) as {
      results: Array<{ index: number; relevance_score: number }>;
    };
    const top = [...data.results]
      .sort((a, b) => b.relevance_score - a.relevance_score)
      .slice(0, 6);
    this.logger.log(
      `Cohere rerank scores top: [${top
        .map((r) => `${r.index}:${r.relevance_score.toFixed(4)}`)
        .join(', ')}] minScore=${config.minScore}`,
    );
    return data.results
      .filter((r) => r.relevance_score >= config.minScore)
      .map<RerankResult<T>>((r) => ({
        index: r.index,
        score: r.relevance_score,
        payload: input.documents[r.index].payload,
      }));
  }

  // ───────────────────────────────────────────────────────────────
  // Jina Rerank API
  //   POST https://api.jina.ai/v1/rerank
  //   docs: https://jina.ai/reranker/
  // ───────────────────────────────────────────────────────────────
  private async rerankJina<T>(
    config: RerankConfig,
    input: RerankInput<T>,
  ): Promise<RerankResult<T>[]> {
    if (!config.apiKey) {
      throw new Error('Jina requiere API key configurada.');
    }
    const url = (config.baseUrl ?? 'https://api.jina.ai').replace(/\/$/, '') + '/v1/rerank';
    const body = {
      model: config.modelo,
      query: input.query,
      documents: input.documents.map((d) => d.text),
      top_n: Math.min(config.topK, input.documents.length),
    };
    const resp = await fetch(url, {
      method: 'POST',
      headers: {
        'Authorization': `Bearer ${config.apiKey}`,
        'Content-Type': 'application/json',
      },
      body: JSON.stringify(body),
    });
    if (!resp.ok) {
      const text = await resp.text().catch(() => '');
      throw new Error(`Jina ${resp.status}: ${text.slice(0, 200)}`);
    }
    const data = (await resp.json()) as {
      results: Array<{ index: number; relevance_score: number }>;
    };
    return data.results
      .filter((r) => r.relevance_score >= config.minScore)
      .map<RerankResult<T>>((r) => ({
        index: r.index,
        score: r.relevance_score,
        payload: input.documents[r.index].payload,
      }));
  }

  // ───────────────────────────────────────────────────────────────
  // Local (Ollama-like). Asume endpoint compatible con Cohere format:
  //   POST {baseUrl}/v1/rerank
  // Si tu Ollama tiene un schema distinto, ajustar acá.
  // ───────────────────────────────────────────────────────────────
  private async rerankLocal<T>(
    config: RerankConfig,
    input: RerankInput<T>,
  ): Promise<RerankResult<T>[]> {
    if (!config.baseUrl) {
      throw new Error('Reranker local requiere base_url configurada.');
    }
    const url = config.baseUrl.replace(/\/$/, '') + '/v1/rerank';
    const body = {
      model: config.modelo,
      query: input.query,
      documents: input.documents.map((d) => d.text),
      top_n: Math.min(config.topK, input.documents.length),
    };
    const resp = await fetch(url, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify(body),
    });
    if (!resp.ok) {
      const text = await resp.text().catch(() => '');
      throw new Error(`Reranker local ${resp.status}: ${text.slice(0, 200)}`);
    }
    const data = (await resp.json()) as {
      results: Array<{ index: number; relevance_score: number }>;
    };
    return data.results
      .filter((r) => r.relevance_score >= config.minScore)
      .map<RerankResult<T>>((r) => ({
        index: r.index,
        score: r.relevance_score,
        payload: input.documents[r.index].payload,
      }));
  }

  // ───────────────────────────────────────────────────────────────
  // Fallback: devuelve los primeros topK en el orden recibido con score 0.
  // Sirve cuando el reranker falla — el caller sigue funcionando con la
  // ordenación vectorial original.
  // ───────────────────────────────────────────────────────────────
  private fallback<T>(input: RerankInput<T>): RerankResult<T>[] {
    return input.documents.map<RerankResult<T>>((d, i) => ({
      index: i,
      score: 0,
      payload: d.payload,
    }));
  }
}