// The "Search in" fold — a run of the rewritten tree, read as the tree the
// visitor built.
//
// `applySearchIn` (lib/searchQuery.ts) turns a "transcripts" leaf that reads
// two or three kinds into an OR over one copy per kind, each copy with an id of
// its own. The run reports under those ids; the builder, the result cards and
// the session all know only the visitor's leaf. This fold files every copy's
// hits and counts back under that leaf, so a chat hit lands in the same
// section of the same video row as the transcript hits, wearing its track.
//
// Pure: it copies what it changes and returns the progress it was given when
// the rewrite changed nothing.

import type { SearchInTree } from "../searchQuery";
import type { LeafState, TreeProgress } from "../searchEval";
import type { LayerHit } from "./leafPipeline";

export function foldSearchIn(
  p: TreeProgress,
  tree: Pick<SearchInTree, "origin" | "unionOf">,
): TreeProgress {
  const { origin, unionOf } = tree;
  if (origin.size === 0) return p;

  const hits = new Map<string, LayerHit[]>();
  for (const [slug, list] of p.hits) {
    let changed = false;
    const out = list.map((h) => {
      const to = origin.get(h.leafId);
      if (to === undefined) return h;
      changed = true;
      return { ...h, leafId: to };
    });
    hits.set(slug, changed ? out : list);
  }

  // Each copy's state, grouped under the leaf it stands for. A copy that has
  // not started yet has no state at all.
  const copies = new Map<string, LeafState[]>();
  const copiesExpected = new Map<string, number>();
  for (const to of origin.values()) {
    copiesExpected.set(to, (copiesExpected.get(to) ?? 0) + 1);
  }
  const leafStates = new Map<string, LeafState>();
  for (const [id, state] of p.leafStates) {
    const to = origin.get(id);
    if (to === undefined) {
      leafStates.set(id, state);
      continue;
    }
    const list = copies.get(to) ?? [];
    list.push(state);
    copies.set(to, list);
  }

  for (const [to, states] of copies) {
    const allStarted = states.length === (copiesExpected.get(to) ?? 0);
    // The OR group's state is the union of what the copies matched, inside the
    // same scope the leaf itself would have read — the count the leaf shows.
    // Until every copy has started it is not: the evaluator reads a copy with
    // no result yet as "everything in scope", so the union would briefly be the
    // whole scope. Until then the largest copy is the count (a lower bound).
    const groupId = unionOf.get(to);
    const union = groupId ? p.groupStates.get(groupId) : undefined;
    let slugCount = 0;
    if (allStarted && union) slugCount = union.slugCount;
    else for (const s of states) slugCount = Math.max(slugCount, s.slugCount);
    let totalHits = 0;
    let processed = 0;
    let totalToProcess = 0;
    let capped = false;
    let cached = true;
    let active = !allStarted;
    for (const s of states) {
      totalHits += s.totalHits;
      processed += s.processed;
      totalToProcess += s.totalToProcess;
      capped ||= s.capped;
      cached &&= s.cached;
      active ||= s.active;
    }
    leafStates.set(to, {
      slugCount,
      totalHits,
      processed,
      totalToProcess,
      capped,
      cached,
      active,
    });
  }

  return { ...p, hits, leafStates };
}