ttrpg_manager/src/features/assistant/useEncounterAdvisor.ts

import { useMemo, useRef, useState } from 'react';
import type { Campaign, Encounter } from '@/lib/schemas';
import { encountersRepo } from '@/lib/db/repositories';
import { newId } from '@/lib/ids';
import { createRng } from '@/lib/rng';
import { rollDice } from '@/lib/dice/notation';
import { addCombatant, removeCombatant } from '@/lib/combat/engine';
import { computeBudget } from '@/lib/combat/budget';
import { loadMonsters, loadPf2e } from '@/lib/compendium';
import { complete } from '@/lib/llm/client';
import { getLlmConfig } from '@/stores/assistantStore';
import { useAssistantStore } from '@/stores/assistantStore';
import { buildCampaignContext, pickCreatureCandidates, type CreatureCandidate } from '@/lib/assistant/context';
import { detectThemes, type CampaignTheme } from '@/lib/assistant/patterns';
import { buildBalancePrompt, balanceSuggestionSchema, type BalanceSuggestion } from '@/lib/assistant/prompts';
import { buildSuggestedEncounter, suggestReinforcements, baseName, type ReinforceMonster } from '@/lib/assistant/encounter';
import { useCharacters } from '@/features/characters/hooks';
import { useNotes, useQuests } from '@/features/world/hooks';
import { useEncounters } from '@/features/combat/hooks';

export type AdvisorState = 'idle' | 'loading' | 'ready' | 'error';

const ORDINAL: Record<string, number> = {
  trivial: 0, easy: 1, low: 1, medium: 2, moderate: 2, hard: 3, severe: 3, deadly: 4, extreme: 4,
};
const TIER_LABEL: Record<Campaign['system'], Record<number, string>> = {
  '5e': { 2: 'Medium', 3: 'Hard', 4: 'Deadly' },
  pf2e: { 2: 'Moderate', 3: 'Severe', 4: 'Extreme' },
};

/** One tier above the current difficulty, floored at Medium/Moderate. */
function nextTarget(system: Campaign['system'], current: string): string {
  const ord = ORDINAL[current] ?? 0;
  const desired = Math.min(4, Math.max(2, ord + 1));
  return TIER_LABEL[system][desired]!;
}

function toCombatant(system: Campaign['system'], m: Record<string, unknown>) {
  const is5e = system === '5e';
  const name = String(m.name);
  const ac = Number(is5e ? m.armor_class : m.ac) || 10;
  const hp = Number(is5e ? m.hit_points : m.hp) || 1;
  const initBonus = is5e ? Math.floor(((Number(m.dexterity) || 10) - 10) / 2) : Number(m.perception) || 0;
  const cr = Number(m.cr);
  const level = Number(m.level);
  return {
    id: newId(), name, kind: 'monster' as const,
    initiative: rollDice('1d20', createRng()).total + initBonus, initBonus,
    ac, hp: { current: hp, max: hp, temp: 0 }, conditions: [], notes: '',
    ...(is5e && Number.isFinite(cr) ? { cr } : {}),
    ...(!is5e && Number.isFinite(level) ? { level } : {}),
  };
}

function countByName(chosen: Record<string, unknown>[]): BalanceSuggestion['add'] {
  const counts = new Map<string, number>();
  for (const m of chosen) {
    const name = String(m.name);
    counts.set(name, (counts.get(name) ?? 0) + 1);
  }
  return [...counts].map(([name, count]) => ({ name, count }));
}

export function useEncounterAdvisor(campaign: Campaign, encounter: Encounter) {
  const characters = useCharacters(campaign.id);
  const notes = useNotes(campaign.id);
  const quests = useQuests(campaign.id);
  const encounters = useEncounters(campaign.id);
  const llmEnabled = useAssistantStore((s) => s.enabled);
  const hasKey = useAssistantStore((s) => !!s.apiKey);

  const ctx = useMemo(
    () => buildCampaignContext({ campaign, characters, encounters, quests, notes }),
    [campaign, characters, encounters, quests, notes],
  );
  const theme = useMemo<CampaignTheme | undefined>(
    () => detectThemes(ctx).find((t) => t.kind === 'encounter-difficulty'),
    [ctx],
  );

  const { hasMonsters, currentDifficulty } = useMemo(() => {
    const monsters = encounter.combatants.filter((c) => c.kind === 'monster').map((c) => ({ cr: c.cr, level: c.level }));
    const diff = monsters.length && ctx.partyLevels.length
      ? computeBudget(campaign.system, ctx.partyLevels, monsters).difficulty
      : undefined;
    return { hasMonsters: monsters.length > 0, currentDifficulty: diff };
  }, [encounter.combatants, ctx.partyLevels, campaign.system]);

  const [state, setState] = useState<AdvisorState>('idle');
  const [suggestion, setSuggestion] = useState<BalanceSuggestion | undefined>();
  const [source, setSource] = useState<'llm' | 'deterministic'>('deterministic');
  const [message, setMessage] = useState<string | null>(null);
  const poolRef = useRef<Record<string, unknown>[]>([]);

  const canUseLlm = llmEnabled && hasKey;

  const run = async () => {
    setState('loading');
    setMessage(null);
    setSuggestion(undefined);
    try {
      const raw = (campaign.system === '5e' ? await loadMonsters() : await loadPf2e('creatures')) as unknown as Record<string, unknown>[];
      poolRef.current = raw;
      const partyLevels = ctx.partyLevels;
      if (partyLevels.length === 0) { setMessage('Add player characters first.'); setState('error'); return; }

      const monsterCombatants = encounter.combatants.filter((c) => c.kind === 'monster');
      const existing: ReinforceMonster[] = monsterCombatants.map((c) => ({ name: c.name, cr: c.cr, level: c.level }));
      const current = computeBudget(campaign.system, partyLevels, existing.map((m) => ({ cr: m.cr, level: m.level }))).difficulty;
      const target = nextTarget(campaign.system, current);

      // Candidates the AI may choose from: the creatures already in the fight
      // (so it can say "add another goblin") plus level-appropriate bestiary picks.
      const existingCands: CreatureCandidate[] = [];
      const seenBase = new Set<string>();
      for (const c of monsterCombatants) {
        const bn = baseName(c.name);
        const rating = campaign.system === '5e' ? c.cr : c.level;
        if (rating === undefined || seenBase.has(bn.toLowerCase())) continue;
        seenBase.add(bn.toLowerCase());
        existingCands.push({ name: bn, rating, ac: c.ac, hp: c.hp.max });
      }
      const poolCands = pickCreatureCandidates(campaign.system, partyLevels, raw, target, 14)
        .filter((c) => !seenBase.has(c.name.toLowerCase()));
      const candidates = [...existingCands, ...poolCands];
      if (candidates.length === 0) { setMessage('No suitable creatures found for this party.'); setState('error'); return; }

      if (canUseLlm) {
        const prompt = buildBalancePrompt(ctx, { difficulty: current, targetDifficulty: target, candidates });
        const res = await complete(getLlmConfig(), { system: prompt.system, user: prompt.user, schema: balanceSuggestionSchema });
        if (res.ok && 'data' in res) {
          const valid = res.data.add.filter((a) => candidates.some((c) => c.name === a.name));
          if (valid.length) {
            setSuggestion({ ...res.data, add: valid });
            setSource('llm');
            setState('ready');
            return;
          }
        } else if (!res.ok) {
          setMessage(`AI unavailable: ${res.message} Showing a deterministic pick.`);
        }
      }

      // Deterministic: reinforce the existing fight; only build fresh if it's empty.
      const typedPool = raw as (Record<string, unknown> & { cr?: number; level?: number })[];
      let plan = suggestReinforcements(campaign.system, partyLevels, existing, typedPool, target);
      if (!plan) {
        const chosen = buildSuggestedEncounter(campaign.system, partyLevels, typedPool, target);
        if (!chosen.length) { setMessage('Could not assemble a balanced suggestion.'); setState('error'); return; }
        plan = { reasoning: `A level-appropriate group to reach ${target} difficulty.`, add: countByName(chosen), targetDifficulty: target };
      }
      setSuggestion(plan);
      setSource('deterministic');
      setState('ready');
    } catch {
      setMessage('Something went wrong building a suggestion.');
      setState('error');
    }
  };

  const apply = async () => {
    if (!suggestion) return;
    await encountersRepo.mutate(encounter.id, (e) => {
      let next = e;
      // Over-tuned fights: take out the suggested monsters (newest copies first).
      for (const r of suggestion.remove ?? []) {
        for (let i = 0; i < r.count; i++) {
          const victim = [...next.combatants].reverse().find(
            (c) => c.kind === 'monster' && baseName(c.name) === r.name,
          );
          if (!victim) break;
          next = removeCombatant(next, victim.id);
        }
      }
      for (const a of suggestion.add) {
        // Prefer cloning a creature already in the fight (handles "add another goblin"
        // and works even for custom/homebrew combatants); otherwise pull from the pool.
        const template = e.combatants.find((c) => c.kind === 'monster' && baseName(c.name) === a.name);
        const make = template
          ? () => ({
              ...template, id: newId(), name: baseName(template.name),
              hp: { current: template.hp.max, max: template.hp.max, temp: 0 },
              conditions: [], initiative: rollDice('1d20', createRng()).total + template.initBonus,
            })
          : (() => {
              const m = poolRef.current.find((p) => String(p.name) === a.name);
              return m ? () => toCombatant(campaign.system, m) : null;
            })();
        if (!make) continue;
        for (let i = 0; i < a.count; i++) next = addCombatant(next, make());
      }
      return next;
    });
    setState('idle');
    setSuggestion(undefined);
  };

  return { theme, hasMonsters, currentDifficulty, state, suggestion, source, message, canUseLlm, run, apply };
}