import type { RerankPort } from "../llm/types";
/**
 * Reranking and position-aware blending.
 * Uses RerankPort to reorder candidates.
 *
 * @module src/pipeline/rerank
 */
import type { ChunkRow, StorePort } from "../store/types";
import type { RequestHydration } from "./hydration";
import type { BlendingTier, FusionCandidate, RerankedCandidate } from "./types";

import {
  assertInferenceActive,
  assertInferenceResult,
} from "../llm/inference-scope";
import {
  buildIntentAwareRerankQuery,
  selectBestChunkForSteering,
} from "./intent";
import { DEFAULT_BLENDING_SCHEDULE } from "./types";

// ─────────────────────────────────────────────────────────────────────────────
// Types
// ─────────────────────────────────────────────────────────────────────────────

export interface RerankOptions {
  /** Max candidates to rerank */
  maxCandidates?: number;
  /** Blending schedule */
  blendingSchedule?: BlendingTier[];
  /** Optional disambiguating context for reranking */
  intent?: string;
  /** Apply bounded auxiliary scoring after fusion normalization, before blend. */
  adjustNormalizedFusionScore?: (
    candidate: FusionCandidate,
    normalizedScore: number
  ) => number;
}

export interface RerankResult {
  candidates: RerankedCandidate[];
  reranked: boolean;
  fallbackReason: "none" | "disabled" | "error";
}

export interface RerankDeps {
  rerankPort: RerankPort | null;
  store: StorePort;
  /** Shared raw chunks only; model inputs are prepared per invocation. */
  hydration?: RequestHydration;
}

// ─────────────────────────────────────────────────────────────────────────────
// Blending
// ─────────────────────────────────────────────────────────────────────────────

/**
 * Get blending weights for a position.
 */
function getBlendingWeights(
  position: number,
  schedule: BlendingTier[]
): { fusionWeight: number; rerankWeight: number } {
  const tier = schedule.find((t) => position <= t.maxRank);
  if (tier) {
    return { fusionWeight: tier.fusionWeight, rerankWeight: tier.rerankWeight };
  }
  // Fallback to last tier
  const last = schedule.at(-1);
  return last
    ? { fusionWeight: last.fusionWeight, rerankWeight: last.rerankWeight }
    : { fusionWeight: 0.5, rerankWeight: 0.5 };
}

/**
 * Blend fusion and rerank scores.
 */
function blend(
  fusionScore: number,
  rerankScore: number,
  position: number,
  schedule: BlendingTier[]
): number {
  const { fusionWeight, rerankWeight } = getBlendingWeights(position, schedule);
  return fusionWeight * fusionScore + rerankWeight * rerankScore;
}

// ─────────────────────────────────────────────────────────────────────────────
// Chunk Text Extraction
// ─────────────────────────────────────────────────────────────────────────────

const MAX_CHUNK_CHARS = 4000;
const PROTECT_BM25_TOP_RANK = 1;

function isProtectedLexicalTopHit(candidate: FusionCandidate): boolean {
  return (
    candidate.bm25Rank === PROTECT_BM25_TOP_RANK &&
    candidate.sources.includes("bm25")
  );
}

/**
 * Fetch chunk texts for reranking.
 */
function rerankOwnerKey(candidate: FusionCandidate): string {
  return `${candidate.mirrorHash}${candidate.documentId === undefined ? "" : `:${candidate.documentId}`}`;
}

async function fetchChunkTexts(
  store: Pick<StorePort, "getChunksBatch">,
  toRerank: FusionCandidate[],
  query: string,
  intent: string | undefined
): Promise<{ texts: string[]; hashToIndex: Map<string, number> }> {
  const uniqueHashes = [
    ...new Set(toRerank.map((candidate) => candidate.mirrorHash)),
  ];
  const chunksBatchResult = await store.getChunksBatch(uniqueHashes);
  const chunksByHash: Map<string, ChunkRow[]> = chunksBatchResult.ok
    ? chunksBatchResult.value
    : new Map();
  const preferredSeqByHash = new Map<string, number>();
  const ownerHashes = new Map(
    toRerank.map((candidate) => [
      rerankOwnerKey(candidate),
      candidate.mirrorHash,
    ])
  );

  for (const candidate of toRerank) {
    assertInferenceActive();
    const existingSeq = preferredSeqByHash.get(rerankOwnerKey(candidate));
    if (existingSeq !== undefined) {
      const existingCandidate = toRerank.find(
        (entry) =>
          rerankOwnerKey(entry) === rerankOwnerKey(candidate) &&
          entry.seq === existingSeq
      );
      if (
        existingCandidate &&
        existingCandidate.fusionScore >= candidate.fusionScore
      ) {
        continue;
      }
    }
    preferredSeqByHash.set(rerankOwnerKey(candidate), candidate.seq);
  }

  const chunkTexts = new Map<string, string>();
  for (const [hash, mirrorHash] of ownerHashes) {
    assertInferenceActive();
    const chunks = chunksByHash.get(mirrorHash);
    const bestChunk = selectBestChunkForSteering(chunks ?? [], query, intent, {
      preferredSeq: preferredSeqByHash.get(hash) ?? null,
      intentWeight: 0.5,
    });
    const text = bestChunk?.text ?? "";
    chunkTexts.set(
      hash,
      text.length > MAX_CHUNK_CHARS
        ? `${text.slice(0, MAX_CHUNK_CHARS)}...`
        : text
    );
  }

  const hashToIndex = new Map<string, number>();
  const texts: string[] = [];
  for (const hash of ownerHashes.keys()) {
    assertInferenceActive();
    hashToIndex.set(hash, texts.length);
    texts.push(chunkTexts.get(hash) ?? "");
  }

  return { texts, hashToIndex };
}

// ─────────────────────────────────────────────────────────────────────────────
// Rerank Implementation
// ─────────────────────────────────────────────────────────────────────────────

/**
 * Rerank candidates using cross-encoder.
 * Falls back to fusion-only if reranking fails.
 */
export async function rerankCandidates(
  deps: RerankDeps,
  query: string,
  candidates: FusionCandidate[],
  options: RerankOptions = {}
): Promise<RerankResult> {
  if (candidates.length === 0) {
    return { candidates: [], reranked: false, fallbackReason: "none" };
  }

  const { rerankPort, store } = deps;
  const maxCandidates = options.maxCandidates ?? 20;
  const schedule = options.blendingSchedule ?? DEFAULT_BLENDING_SCHEDULE;

  // Normalize fusion scores to 0-1 range across ALL candidates for stability.
  const fusionScoresAll = candidates.map((c) => c.fusionScore);
  const minFusionAll = Math.min(...fusionScoresAll);
  const maxFusionAll = Math.max(...fusionScoresAll);
  const fusionRangeAll = maxFusionAll - minFusionAll;

  const normalizeFusionScore = (score: number): number => {
    if (fusionRangeAll < 1e-9) {
      return 1;
    }
    const v = (score - minFusionAll) / fusionRangeAll;
    return Math.max(0, Math.min(1, v));
  };
  const adjustedFusionScore = (candidate: FusionCandidate): number => {
    const normalized = normalizeFusionScore(candidate.fusionScore);
    const adjusted = options.adjustNormalizedFusionScore?.(
      candidate,
      normalized
    );
    return adjusted === undefined
      ? normalized
      : Math.max(0, Math.min(1, adjusted));
  };
  const sortAdjustedCandidates = (
    adjustedCandidates: RerankedCandidate[]
  ): RerankedCandidate[] => {
    if (!options.adjustNormalizedFusionScore) {
      return adjustedCandidates;
    }
    return adjustedCandidates.sort(
      (left, right) =>
        right.blendedScore - left.blendedScore ||
        `${left.mirrorHash}:${left.seq}`.localeCompare(
          `${right.mirrorHash}:${right.seq}`
        )
    );
  };

  // No reranker: return candidates with normalized fusion scores
  if (!rerankPort) {
    return {
      candidates: sortAdjustedCandidates(
        candidates.map((c) => ({
          ...c,
          rerankScore: null,
          blendedScore: adjustedFusionScore(c),
        }))
      ),
      reranked: false,
      fallbackReason: "disabled",
    };
  }

  const toRerank = candidates.slice(0, maxCandidates);
  const remaining = candidates.slice(maxCandidates);

  // Extract best chunk per document for efficient reranking
  const { texts, hashToIndex } = await fetchChunkTexts(
    deps.hydration ?? store,
    toRerank,
    query,
    options.intent
  );

  const uniqueTexts: string[] = [];
  const docIndexToUniqueIndex = new Map<number, number>();
  const uniqueIndexToDocIndices = new Map<number, number[]>();
  const textToUniqueIndex = new Map<string, number>();

  for (const [docIndex, text] of texts.entries()) {
    assertInferenceActive();
    const existingIndex = textToUniqueIndex.get(text);
    if (existingIndex !== undefined) {
      docIndexToUniqueIndex.set(docIndex, existingIndex);
      const mapped = uniqueIndexToDocIndices.get(existingIndex) ?? [];
      mapped.push(docIndex);
      uniqueIndexToDocIndices.set(existingIndex, mapped);
      continue;
    }

    const uniqueIndex = uniqueTexts.length;
    uniqueTexts.push(text);
    textToUniqueIndex.set(text, uniqueIndex);
    docIndexToUniqueIndex.set(docIndex, uniqueIndex);
    uniqueIndexToDocIndices.set(uniqueIndex, [docIndex]);
  }

  // Run reranking on best chunks (much faster than full docs)
  const rerankResult = await rerankPort.rerank(
    buildIntentAwareRerankQuery(query, options.intent),
    uniqueTexts
  );

  assertInferenceResult(rerankResult);
  if (!rerankResult.ok) {
    return {
      candidates: sortAdjustedCandidates(
        candidates.map((c) => ({
          ...c,
          rerankScore: null,
          blendedScore: adjustedFusionScore(c),
        }))
      ),
      reranked: false,
      fallbackReason: "error",
    };
  }

  // Normalize rerank scores using min-max
  const scoreByDocIndex = new Map<number, number>();
  for (const score of rerankResult.value) {
    assertInferenceActive();
    const docIndices = uniqueIndexToDocIndices.get(score.index) ?? [];
    for (const docIndex of docIndices) {
      assertInferenceActive();
      scoreByDocIndex.set(docIndex, score.score);
    }
  }
  const rerankScores = rerankResult.value.map((s) => s.score);
  const minRerank = Math.min(...rerankScores);
  const maxRerank = Math.max(...rerankScores);
  const rerankRange = maxRerank - minRerank;

  const normalizeRerankScore = (score: number): number => {
    if (rerankRange < 1e-9) {
      return 1;
    }
    return (score - minRerank) / rerankRange;
  };

  // Build reranked candidates with blended scores
  const rerankedCandidates: RerankedCandidate[] = toRerank.map((c, i) => {
    const docIndex = hashToIndex.get(rerankOwnerKey(c)) ?? -1;
    const rerankScore = scoreByDocIndex.get(docIndex) ?? null;
    const normalizedRerankScore =
      rerankScore !== null ? normalizeRerankScore(rerankScore) : null;

    const position = i + 1;
    const normalizedFusion = adjustedFusionScore(c);
    const blendedScore =
      normalizedRerankScore !== null
        ? blend(normalizedFusion, normalizedRerankScore, position, schedule)
        : normalizedFusion;

    return { ...c, rerankScore: normalizedRerankScore, blendedScore };
  });

  // Add remaining candidates with penalty
  let allCandidates: RerankedCandidate[] = [
    ...rerankedCandidates,
    ...remaining.map((c) => ({
      ...c,
      rerankScore: null,
      blendedScore: Math.max(0, Math.min(1, adjustedFusionScore(c) * 0.5)),
    })),
  ];

  // Sort by blended score with deterministic tie-breaking
  allCandidates.sort((a, b) => {
    const scoreDiff = b.blendedScore - a.blendedScore;
    if (Math.abs(scoreDiff) > 1e-9) {
      return scoreDiff;
    }
    return `${a.mirrorHash}:${a.seq}`.localeCompare(`${b.mirrorHash}:${b.seq}`);
  });

  // Guardrail: keep strong original lexical #1 at the top.
  // This avoids rerank-only demotions on clear exact-hit queries.
  const protectedTopHit = allCandidates.find(isProtectedLexicalTopHit);
  if (protectedTopHit && allCandidates[0] !== protectedTopHit) {
    allCandidates = [
      protectedTopHit,
      ...allCandidates.filter((candidate) => candidate !== protectedTopHit),
    ];
  }

  return { candidates: allCandidates, reranked: true, fallbackReason: "none" };
}