// ---------------------------------------------------------------------------
// Memory v2 — Per-turn activation update
// ---------------------------------------------------------------------------
//
// Implements the activation formula from §4 of the design doc plus an
// additive cross-encoder rerank boost on the unified top-K-by-A_o pool:
//
//   A_o(n, t+1) = d · A(n, t)
//               + c_user      · sim(User_{t+1},  n)
//               + c_assistant · sim(Assistant_t, n)
//               + c_now       · sim(NOW.md,      n)
//               + c_user      · α · r_norm(User_{t+1},  n)   [n ∈ topK]
//               + c_assistant · α · r_norm(Assistant_t, n)   [n ∈ topK]
//
// The spread step that turns A_o into the final A(n, t+1) lives in the
// substrate (`substrate/spread.ts` — `spreadActivation`) because the `memory`
// recall source also spreads activation under v3. Callers import it from there
// directly; this module deliberately does not re-export it, so the symbol has
// one import path.
//
// Candidate selection (§6) keeps the per-turn cost linear in the size of the
// active set rather than the entire concept-page collection. The candidate
// set is the union of:
//   - Slugs whose prior activation exceeds `epsilon` (the persisted state).
//   - The top-50 by ANN hybrid query against `concat(user, assistant, now)` —
//     a single batched call to `hybridQueryConceptPages` with no slug
//     restriction. Pages outside the candidate set decay via `d · A(n, t)`
//     for the next turn and drop below `epsilon` if no longer relevant.

import type { AssistantConfig } from "../../../../config/types.js";
import { isEmbeddingDimensionAvailable } from "../../../../persistence/embeddings/embedding-backend.js";
import { applyCorrectionIfCalibrated } from "../anisotropy.js";
import { embedWithBackend } from "../embeddings.js";
import { hybridQueryConceptPages } from "../substrate/qdrant.js";
import { simBatch } from "../substrate/sim.js";
import { generateBm25QueryEmbedding } from "../substrate/sparse-bm25.js";
import type { ActivationState, EverInjectedEntry } from "../substrate/types.js";
import { clampUnitInterval } from "../validation.js";
import { rerankCandidates } from "./reranker.js";

/**
 * Sentinel passed to Qdrant when `config.memory.v2.ann_candidate_limit` is
 * `null` (unlimited). Qdrant's query API requires an explicit numeric
 * `limit`, so unlimited is represented as a number large enough that any
 * realistic concept-page collection is returned in full.
 *
 * Why not `Number.MAX_SAFE_INTEGER`: Qdrant's sparse-vector `SearchContext`
 * pre-allocates `limit * 16` bytes per query, so passing `MAX_SAFE_INTEGER`
 * triggers a ~144 PB allocation and SIGABRTs the Qdrant process. 1_000_000
 * is ~16 MB of pre-allocation in Qdrant — generous headroom over realistic
 * concept-page counts (low thousands today) while staying well clear of
 * the OOM cliff. Bump explicitly via `ann_candidate_limit` if you ever
 * outgrow it.
 */
const UNLIMITED_ANN_CANDIDATE_LIMIT = 1_000_000;

// ---------------------------------------------------------------------------
// Candidate selection
// ---------------------------------------------------------------------------

interface SelectCandidatesParams {
  /**
   * Prior-turn activation snapshot. Slugs with activation strictly greater
   * than `config.memory.v2.epsilon` are carried forward as candidates so the
   * decay term `d · A(n, t)` continues to influence them next turn.
   */
  priorState: ActivationState | null;
  /** User message text for this turn. */
  userText: string;
  /** Assistant message text from the prior turn (empty string at conv start). */
  assistantText: string;
  /** NOW context string (essentials/threads/recent or NOW.md). */
  nowText: string;
  config: AssistantConfig;
  signal?: AbortSignal;
}

interface SelectCandidatesResult {
  /** Union of `fromPrior` and `fromAnn` — the per-turn candidate set. */
  candidates: Set<string>;
  /** Slugs carried forward from `priorState` because their activation > epsilon. */
  fromPrior: Set<string>;
  /** Slugs surfaced by the unrestricted ANN top-50 against the joined turn text. */
  fromAnn: Set<string>;
}

/**
 * Build the per-turn candidate set: the union of slugs in the prior state
 * (above epsilon) and the top-50 ANN hits against the concatenated turn
 * text. The ANN call runs un-restricted (no slug filter) so it can surface
 * pages outside the active set.
 *
 * Returns the union plus the two source sets separately so downstream
 * telemetry can attribute each candidate to `prior_state`, `ann_top50`, or
 * both. A slug present in both sources appears in `fromPrior ∩ fromAnn`.
 *
 * Empty candidate sets are valid and propagate downstream — both
 * `computeOwnActivation` and `spreadActivation` short-circuit on them.
 */
export async function selectCandidates(
  params: SelectCandidatesParams,
): Promise<SelectCandidatesResult> {
  const { priorState, userText, assistantText, nowText, config, signal } =
    params;

  const fromPrior = new Set<string>();
  const fromAnn = new Set<string>();

  // (1) Carry forward prior-state slugs above epsilon.
  if (priorState) {
    const epsilon = config.memory.v2.epsilon;
    for (const [slug, activation] of Object.entries(priorState.state)) {
      if (activation > epsilon) {
        fromPrior.add(slug);
      }
    }
  }

  // (2) ANN top-50 against the concatenated turn text. Pure whitespace joins
  // (no separators) keep the embedding behavior aligned with how callers
  // would naturally read the three texts together. Whitespace-only channels
  // contribute no semantic content, so trim before deciding whether to embed.
  const annQueryText = [userText, assistantText, nowText]
    .filter((s) => s.trim().length > 0)
    .join("\n");

  if (annQueryText.trim().length > 0) {
    // Dense short-circuit: when the reachable backend can't produce vectors of
    // the committed collection dimension (degraded backend or dimension
    // mismatch), skip the embed round-trip that `embedWithBackend` would only
    // reject on its dimension assertion every turn. `hybridQueryConceptPages`
    // treats an empty dense vector as sparse-only, so the ANN candidate scan
    // still runs on BM25 — the dense channel just contributes nothing.
    let dense: number[] = [];
    if (await isEmbeddingDimensionAvailable(config)) {
      throwIfAborted(signal);
      const denseResult = await embedWithBackend(config, [annQueryText], {
        signal,
      });
      dense = await applyCorrectionIfCalibrated(
        denseResult.vectors[0],
        denseResult.provider,
        denseResult.model,
      );
    }
    throwIfAborted(signal);
    const sparse = generateBm25QueryEmbedding(annQueryText);
    const limit =
      config.memory.v2.ann_candidate_limit ?? UNLIMITED_ANN_CANDIDATE_LIMIT;
    const hits = await hybridQueryConceptPages(dense, sparse, limit);
    for (const hit of hits) {
      fromAnn.add(hit.slug);
    }
  }

  const candidates = new Set<string>([...fromPrior, ...fromAnn]);

  return { candidates, fromPrior, fromAnn };
}

// ---------------------------------------------------------------------------
// Own activation
// ---------------------------------------------------------------------------

interface ComputeOwnActivationParams {
  candidates: ReadonlySet<string>;
  priorState: ActivationState | null;
  userText: string;
  assistantText: string;
  nowText: string;
  config: AssistantConfig;
  signal?: AbortSignal;
}

/**
 * Per-slug breakdown of the own-activation inputs, captured before any
 * coefficient weighting is applied. Surfaced for telemetry / inspector views
 * so the UI can show how each term contributed to the final value.
 */
interface OwnActivationBreakdown {
  /** `d * prev(slug)` — the decayed prior-turn activation contribution. */
  priorContribution: number;
  /** Raw fused `sim(user, slug)`, before `c_user` weighting. */
  simUser: number;
  /** Raw fused `sim(assistant, slug)`, before `c_assistant` weighting. */
  simAssistant: number;
  /** Raw fused `sim(now, slug)`, before `c_now` weighting. */
  simNow: number;
  /** Rerank delta `α · r_norm_u`; 0 outside the top-K pool. Applied to `A_o` weighted by `c_user`. */
  simUserRerankBoost: number;
  /** Rerank delta `α · r_norm_a`; 0 outside the top-K pool. Applied to `A_o` weighted by `c_assistant`. NOW skips rerank. */
  simAssistantRerankBoost: number;
  /** True when this slug was in the unified top-K rerank pool. Lets the inspector distinguish "cross-encoder normalised to 0" from "rerank skipped this slug." */
  inRerankPool: boolean;
}

interface ComputeOwnActivationResult {
  /** Final clamped own-activation value per slug. */
  activation: Map<string, number>;
  /** Per-slug breakdown of the inputs that fed into `activation`. */
  breakdown: Map<string, OwnActivationBreakdown>;
}

/**
 * Apply the own-activation formula
 *   A_o(n) = d · prev(n)
 *          + c_user · sim_u + c_assistant · sim_a + c_now · sim_n
 *          + c_user · α · r_norm_u + c_assistant · α · r_norm_a
 * over the candidate set, where the rerank terms only fire for slugs that
 * land in the unified top-K window. The pool is ranked by the rerank-eligible
 * channels alone (`c_user · sim_u + c_assistant · sim_a`) so prior- or
 * NOW-heavy slugs — which can't gain from rerank — don't starve out
 * genuinely user/assistant-relevant slugs. Returns a sparse map keyed by
 * slug; slugs whose computed value rounds to 0 are still included so callers
 * can see the candidate set explicitly. Also returns a per-slug breakdown of
 * the raw inputs (decayed prior + raw sims + rerank deltas) so callers can
 * render contribution diagnostics without re-running the math.
 *
 * The three `simBatch` calls run concurrently — they hit independent named
 * vectors and embed independent query texts. Cross-encoder rerank then runs
 * once on the unified top-K so an entry strong in both channels can't
 * double-boost itself past entries that only land in one channel.
 */
export async function computeOwnActivation(
  params: ComputeOwnActivationParams,
): Promise<ComputeOwnActivationResult> {
  const {
    candidates,
    priorState,
    userText,
    assistantText,
    nowText,
    config,
    signal,
  } = params;

  const activation = new Map<string, number>();
  const breakdown = new Map<string, OwnActivationBreakdown>();
  if (candidates.size === 0) {
    return { activation, breakdown };
  }

  const { d, c_user, c_assistant, c_now } = config.memory.v2;
  const slugList = [...candidates];

  // NOW context is structured (timestamps, current focus) — outside the
  // cross-encoder's training distribution, so it never participates in rerank.
  const [simUser, simAssistant, simNow] = await Promise.all([
    simBatch(userText, slugList, config, { signal }),
    simBatch(assistantText, slugList, config, { signal }),
    simBatch(nowText, slugList, config, { signal }),
  ]);

  interface SlugInputs {
    slug: string;
    priorContribution: number;
    simU: number;
    simA: number;
    simN: number;
    /** Pre-rerank A_o; full sum used for the final activation value. */
    preRerank: number;
    /**
     * Ranking signal for the unified rerank pool — only the channels that
     * actually participate in rerank (user + assistant). Excluding
     * `priorContribution` and `c_now * simN` prevents prior- or NOW-heavy
     * slugs from consuming the rerank budget despite being ineligible for
     * cross-encoder gains.
     */
    rerankPoolScore: number;
  }
  const inputs: SlugInputs[] = slugList.map((slug) => {
    const prev = priorState?.state[slug] ?? 0;
    const simU = simUser.get(slug) ?? 0;
    const simA = simAssistant.get(slug) ?? 0;
    const simN = simNow.get(slug) ?? 0;
    const priorContribution = d * prev;
    const rerankPoolScore = c_user * simU + c_assistant * simA;
    return {
      slug,
      priorContribution,
      simU,
      simA,
      simN,
      preRerank: priorContribution + rerankPoolScore + c_now * simN,
      rerankPoolScore,
    };
  });

  // Unified top-K by rerank-eligible signal only. Both channels rerank against
  // the **same** slug set, so a slug strong on user can't crowd out one strong
  // on assistant by virtue of appearing in both per-channel top-Ks. Both
  // channel queries ride in a single `rerankCandidates` call so the worker
  // tokenizes and forward-passes them together — half the per-call overhead
  // of two serialised round-trips.
  let userRerankBoost: ReadonlyMap<string, number> = new Map();
  let assistantRerankBoost: ReadonlyMap<string, number> = new Map();
  let inPoolSet: ReadonlySet<string> = new Set();
  const rerankCfg = config.memory.v2.rerank;
  if (rerankCfg?.enabled) {
    throwIfAborted(signal);
    const topSlugs = inputs
      .slice()
      .sort((a, b) => b.rerankPoolScore - a.rerankPoolScore)
      .slice(0, rerankCfg.top_k)
      .map((e) => e.slug);
    if (topSlugs.length > 0) {
      const [userScores, assistantScores] = await rerankCandidates(
        [userText, assistantText],
        topSlugs,
        config,
      );
      throwIfAborted(signal);
      // Build the pool from slugs the cross-encoder actually scored, so a
      // backend failure (which yields empty maps) doesn't mislabel candidates
      // as `inRerankPool` in the inspector.
      inPoolSet = new Set([...userScores.keys(), ...assistantScores.keys()]);
      userRerankBoost = normalizeRerankScores(userScores, rerankCfg.alpha);
      assistantRerankBoost = normalizeRerankScores(
        assistantScores,
        rerankCfg.alpha,
      );
    }
  }

  for (const e of inputs) {
    const boostU = userRerankBoost.get(e.slug) ?? 0;
    const boostA = assistantRerankBoost.get(e.slug) ?? 0;
    activation.set(
      e.slug,
      clampUnitInterval(e.preRerank + c_user * boostU + c_assistant * boostA),
    );
    breakdown.set(e.slug, {
      priorContribution: e.priorContribution,
      simUser: e.simU,
      simAssistant: e.simA,
      simNow: e.simN,
      simUserRerankBoost: boostU,
      simAssistantRerankBoost: boostA,
      inRerankPool: inPoolSet.has(e.slug),
    });
  }

  return { activation, breakdown };
}

/**
 * Per-batch normalisation: divide raw cross-encoder scores by the channel's
 * own max and return `alpha · r_norm` per slug. Empty input or all-zero
 * scores yield an empty Map so the channel contributes 0 boost.
 */
function normalizeRerankScores(
  rawScores: ReadonlyMap<string, number>,
  alpha: number,
): Map<string, number> {
  const out = new Map<string, number>();
  if (rawScores.size === 0) {
    return out;
  }
  let maxScore = 0;
  for (const v of rawScores.values()) {
    if (v > maxScore) {
      maxScore = v;
    }
  }
  if (maxScore === 0) {
    return out;
  }
  for (const [slug, raw] of rawScores) {
    out.set(slug, alpha * (raw / maxScore));
  }
  return out;
}

function throwIfAborted(signal: AbortSignal | undefined): void {
  if (signal?.aborted) {
    throw new DOMException("Aborted", "AbortError");
  }
}

// ---------------------------------------------------------------------------
// Injection selection
// ---------------------------------------------------------------------------

interface SelectInjectionsParams {
  /** Final activation map after spread. */
  A: ReadonlyMap<string, number>;
  /** Slugs already attached to a prior user message (with their turn). */
  priorEverInjected: readonly EverInjectedEntry[];
  /** Cap on the per-turn injection slate, e.g. `config.memory.v2.top_k`. */
  topK: number;
}

interface SelectInjectionsResult {
  /** Top-K slugs by activation (descending), used for the cached top-now view. */
  topNow: string[];
  /**
   * Slugs in `topNow` that have not yet been attached to any prior user
   * message — the new injections to render on the current user message.
   */
  toInject: string[];
}

/**
 * Pick the top-K slugs by activation (descending; stable on ties via slug
 * lexicographic order) and subtract slugs already in `priorEverInjected` to
 * yield the per-turn injection delta. Empty activation map → empty results.
 */
export function selectInjections(
  params: SelectInjectionsParams,
): SelectInjectionsResult {
  const { A, priorEverInjected, topK } = params;
  if (A.size === 0 || topK <= 0) {
    return { topNow: [], toInject: [] };
  }

  const ranked = [...A.entries()].sort(([slugA, valA], [slugB, valB]) => {
    if (valB !== valA) {
      return valB - valA;
    } // higher activation first
    return slugA < slugB ? -1 : slugA > slugB ? 1 : 0; // stable tie-break
  });

  const topNow = ranked.slice(0, topK).map(([slug]) => slug);
  const everSet = new Set(priorEverInjected.map((entry) => entry.slug));
  const toInject = topNow.filter((slug) => !everSet.has(slug));

  return { topNow, toInject };
}