ttrpg_manager/src/lib/assistant/encounter.ts

import type { SystemId } from '@/lib/rules';
import type { Rng } from '@/lib/rng';
import { createRng } from '@/lib/rng';
import { computeBudget } from '@/lib/combat/budget';

export interface PoolMonster {
  cr?: number | undefined;
  level?: number | undefined;
}

function shuffle<T>(arr: T[], rng: Rng): T[] {
  const a = [...arr];
  for (let i = a.length - 1; i > 0; i--) {
    const j = rng.int(0, i);
    [a[i], a[j]] = [a[j]!, a[i]!];
  }
  return a;
}

/**
 * Greedily assemble monsters from a pool to reach (at least) the target
 * difficulty for the party, choosing level/CR-appropriate creatures. Pure given
 * an Rng. Returns the chosen pool entries (may repeat a creature).
 */
export function buildSuggestedEncounter<T extends PoolMonster>(
  system: SystemId,
  partyLevels: number[],
  pool: T[],
  difficulty: string,
  rng: Rng = createRng(),
  maxMonsters = 8,
): T[] {
  const avg = partyLevels.length ? Math.round(partyLevels.reduce((a, b) => a + b, 0) / partyLevels.length) : 1;
  const thresholds = computeBudget(system, partyLevels, []).thresholds;
  const target = thresholds.find((t) => t.label.toLowerCase() === difficulty.toLowerCase())?.value ?? 0;
  if (target <= 0) return [];

  const eligible = pool.filter((m) =>
    system === '5e'
      ? m.cr !== undefined && m.cr <= avg + 1 && m.cr >= Math.max(0, avg - 6)
      : m.level !== undefined && m.level <= avg + 2 && m.level >= avg - 4,
  );
  if (eligible.length === 0) return [];

  const order = shuffle(eligible, rng);
  const chosen: T[] = [];
  for (let i = 0; chosen.length < maxMonsters; i++) {
    chosen.push(order[i % order.length]!);
    if (computeBudget(system, partyLevels, chosen).ratingXp >= target) break;
  }
  return chosen;
}

// ---------------- Reinforcing an existing encounter ----------------

export interface ReinforceMonster {
  name: string;
  cr?: number | undefined;
  level?: number | undefined;
}

export interface ReinforcePlan {
  reasoning: string;
  add: { name: string; count: number }[];
  /** Monsters to take out when the fight is already at/above the target (over-tuned). */
  remove?: { name: string; count: number }[];
  targetDifficulty: string;
}

/** Strip a combat auto-number suffix, e.g. "Goblin 2" → "Goblin". */
export function baseName(name: string): string {
  return name.replace(/\s+\d+$/, '').trim();
}

function tokens(name: string): string[] {
  return baseName(name).toLowerCase().split(/\s+/).filter((t) => t.length >= 3);
}

function ratingOf(system: SystemId, m: { cr?: number | undefined; level?: number | undefined }): number | undefined {
  return system === '5e' ? m.cr : m.level;
}

/**
 * Suggest how to raise an EXISTING encounter to a target difficulty by reinforcing
 * what's already there — more of the same creatures, or a thematically related
 * stronger creature from the bestiary — rather than proposing a fresh encounter.
 * Accounts for the 5e encounter multiplier (which grows with monster count).
 * Returns null when there are no existing monsters to build on.
 */
export function suggestReinforcements(
  system: SystemId,
  partyLevels: number[],
  existing: ReinforceMonster[],
  pool: (Record<string, unknown> & { cr?: number; level?: number })[],
  targetDifficulty: string,
  maxAdds = 10,
): ReinforcePlan | null {
  if (existing.length === 0) return null;
  const thresholds = computeBudget(system, partyLevels, []).thresholds;
  const target = thresholds.find((t) => t.label.toLowerCase() === targetDifficulty.toLowerCase())?.value ?? 0;
  if (target <= 0) return null;

  // If the fight is already AT or ABOVE the target, it's tuned correctly or
  // over-tuned — adding more would only make it harder. Suggest softening it
  // instead of reinforcing it.
  const currentXp = computeBudget(
    system,
    partyLevels,
    existing.map((m) => ({ cr: m.cr, level: m.level })),
  ).ratingXp;
  if (currentXp >= target) {
    return easeReinforcements(system, partyLevels, existing, target, targetDifficulty);
  }

  const minCountToReach = (rating: number): number | null => {
    for (let n = 1; n <= maxAdds; n++) {
      const extra = Array.from({ length: n }, () => (system === '5e' ? { cr: rating } : { level: rating }));
      if (computeBudget(system, partyLevels, [...existing, ...extra]).ratingXp >= target) return n;
    }
    return null;
  };

  // Candidate reinforcements: distinct existing creatures, plus thematic upgrades
  // from the pool (sharing a name token, level/CR-appropriate, not already present).
  type Cand = { name: string; rating: number; thematic: boolean };
  const cands: Cand[] = [];
  const seen = new Set<string>();
  for (const m of existing) {
    const bn = baseName(m.name);
    const r = ratingOf(system, m);
    if (r === undefined || seen.has(bn.toLowerCase())) continue;
    seen.add(bn.toLowerCase());
    cands.push({ name: bn, rating: r, thematic: false });
  }

  const avg = partyLevels.length ? Math.round(partyLevels.reduce((a, b) => a + b, 0) / partyLevels.length) : 1;
  const eligible = (r: number) => (system === '5e' ? r <= avg + 1 && r >= Math.max(0, avg - 6) : r <= avg + 2 && r >= avg - 4);
  const existingTokens = new Set(existing.flatMap((m) => tokens(m.name)));
  for (const p of pool) {
    const r = ratingOf(system, p);
    const name = typeof p.name === 'string' ? p.name : undefined;
    if (r === undefined || !name || !eligible(r) || seen.has(name.toLowerCase())) continue;
    if (tokens(name).some((t) => existingTokens.has(t))) {
      cands.push({ name, rating: r, thematic: true });
      seen.add(name.toLowerCase());
    }
  }

  // Pick the plan that reaches the target with the fewest added bodies; on ties,
  // prefer the higher-rated (more "boss-like") and thematic option.
  let best: { name: string; count: number; rating: number; thematic: boolean } | null = null;
  for (const c of cands) {
    const n = minCountToReach(c.rating);
    if (n === null) continue;
    const better =
      !best || n < best.count ||
      (n === best.count && (c.rating > best.rating || (c.rating === best.rating && c.thematic && !best.thematic)));
    if (better) best = { name: c.name, count: n, rating: c.rating, thematic: c.thematic };
  }

  // If nothing reaches within the cap, pile on the strongest existing creature.
  if (!best) {
    const strongest = cands.reduce((a, b) => (b.rating > a.rating ? b : a));
    best = { name: strongest.name, count: maxAdds, rating: strongest.rating, thematic: strongest.thematic };
  }

  const more = best.thematic ? '' : ' more';
  const reasoning =
    best.count === 1
      ? `Add 1 ${best.name} to push this fight toward ${targetDifficulty}.`
      : `Add ${best.count}${more} ${best.name} to push this fight toward ${targetDifficulty}.`;
  return { reasoning, add: [{ name: best.name, count: best.count }], targetDifficulty };
}

/**
 * The fight already meets or exceeds the target difficulty. Suggest taking out the
 * fewest bodies that drops it back UNDER the target so it's no harder than intended.
 * Removing creatures also shrinks the 5e encounter multiplier, so the budget is
 * recomputed after each removal rather than subtracted linearly. We never empty the
 * encounter: if even leaving a single creature stays at/above target, we say so and
 * point at swapping for something weaker.
 */
function easeReinforcements(
  system: SystemId,
  partyLevels: number[],
  existing: ReinforceMonster[],
  target: number,
  targetDifficulty: string,
): ReinforcePlan {
  // Strongest first: dropping the heaviest hitters eases the fight with the fewest
  // bodies removed. Stable on equal ratings so ties resolve deterministically.
  const order = existing
    .map((m, i) => ({ m, i, rating: ratingOf(system, m) ?? -Infinity }))
    .sort((a, b) => b.rating - a.rating || a.i - b.i);

  const removedCounts = new Map<string, number>();
  const keep = [...order];
  // Take from the front (strongest) but always leave at least one creature behind.
  while (keep.length > 1) {
    const dropped = keep.shift()!;
    const bn = baseName(dropped.m.name);
    removedCounts.set(bn, (removedCounts.get(bn) ?? 0) + 1);
    const remaining = keep.map((e) => ({ cr: e.m.cr, level: e.m.level }));
    if (computeBudget(system, partyLevels, remaining).ratingXp < target) break;
  }

  const remove = [...removedCounts].map(([name, count]) => ({ name, count }));
  const totalRemoved = remove.reduce((s, r) => s + r.count, 0);

  if (totalRemoved === 0) {
    // Couldn't trim anything (a single over-budget creature). Suggest a swap instead.
    const solo = baseName(existing[0]!.name);
    return {
      reasoning:
        `This fight is already at or above ${targetDifficulty}. ` +
        `Swap the ${solo} for a weaker creature to soften it.`,
      add: [],
      targetDifficulty,
    };
  }

  const phrase = remove
    .map((r) => `${r.count} ${r.name}`)
    .join(' and ');
  const reasoning =
    `This fight is already at or above ${targetDifficulty}. ` +
    `Remove ${phrase} to ease it back toward ${targetDifficulty}.`;
  return { reasoning, add: [], remove, targetDifficulty };
}