Codewalkers/packages/web/src/lib/invalidation.ts

import type { QueryClient } from "@tanstack/react-query";
import { MutationCache } from "@tanstack/react-query";
import type { AnyQueryProcedure, AnyMutationProcedure } from "@trpc/server";
import type { AppRouter } from "@codewalk-district/shared";

// Strip the [key: string] index signature from RouterRecord so keyof yields
// only the literal procedure names, not `string`.
type RemoveIndexSignature<T> = {
  [K in keyof T as string extends K ? never : K]: T[K];
};

type Procedures = RemoveIndexSignature<AppRouter>;

type MutationName = {
  [K in keyof Procedures]: Procedures[K] extends AnyMutationProcedure ? K : never;
}[keyof Procedures] & string;

type QueryName = {
  [K in keyof Procedures]: Procedures[K] extends AnyQueryProcedure ? K : never;
}[keyof Procedures] & string;

/**
 * Centralized invalidation map.
 *
 * Maps each tRPC mutation name to the query keys that should be invalidated
 * when that mutation succeeds. This eliminates scattered `utils.listX.invalidate()`
 * calls across every component.
 *
 * tRPC React Query encodes keys as arrays: the first element is a tuple like
 * ["listAgents"], and the mutation key follows the same pattern. We match on
 * the procedure name (the last segment of the tRPC path).
 */

const INVALIDATION_MAP: Partial<Record<MutationName, QueryName[]>> = {
  // --- Agents ---
  stopAgent: ["listAgents", "listWaitingAgents", "listMessages"],
  deleteAgent: ["listAgents"],
  dismissAgent: ["listAgents"],
  resumeAgent: ["listAgents", "listWaitingAgents", "listMessages"],
  respondToMessage: ["listWaitingAgents", "listMessages"],

  // --- Architect spawns ---
  spawnArchitectRefine: ["listAgents"],
  spawnArchitectDiscuss: ["listAgents"],
  spawnArchitectBreakdown: ["listAgents"],
  spawnArchitectDecompose: ["listAgents", "listInitiativeTasks"],

  // --- Initiatives ---
  createInitiative: ["listInitiatives"],
  updateInitiative: ["listInitiatives", "getInitiative"],
  updateInitiativeProjects: ["getInitiative"],

  // --- Phases ---
  createPhase: ["listPhases", "listInitiativePhaseDependencies"],
  deletePhase: ["listPhases", "listInitiativeTasks", "listInitiativePhaseDependencies"],
  updatePhase: ["listPhases", "getPhase"],
  approvePhase: ["listPhases", "listInitiativeTasks"],
  queuePhase: ["listPhases"],
  createPhaseDependency: ["getPhaseDependencies", "listInitiativePhaseDependencies"],
  removePhaseDependency: ["getPhaseDependencies", "listInitiativePhaseDependencies"],

  // --- Tasks ---
  createPhaseTask: ["listPhaseTasks", "listInitiativeTasks", "listTasks"],
  createInitiativeTask: ["listTasks", "listInitiativeTasks"],
  createChildTasks: ["listTasks", "listInitiativeTasks", "listPhaseTasks"],
  queueTask: ["listTasks", "listInitiativeTasks", "listPhaseTasks"],
  approveTask: ["listTasks", "listInitiativeTasks", "listPhaseTasks", "listPendingApprovals"],

  // --- Change Sets ---
  revertChangeSet: ["listPhases", "listPhaseTasks", "listInitiativeTasks", "listPages", "getPage", "listChangeSets", "getRootPage"],

  // --- Pages ---
  updatePage: ["listPages", "getPage", "getRootPage"],
  createPage: ["listPages", "getRootPage"],
  deletePage: ["listPages", "getRootPage"],

  // --- Projects ---
  registerProject: ["listProjects"],

  // --- Accounts ---
  addAccount: ["listAccounts"],
  removeAccount: ["listAccounts"],
  refreshAccounts: ["listAccounts"],
  markAccountExhausted: ["listAccounts"],
};

/**
 * Extract the tRPC procedure name from a mutation key.
 *
 * tRPC v11 React Query keys look like: [["procedureName"], { type: "mutation" }]
 * We want just "procedureName".
 */
function extractProcedureName(mutationKey: unknown): MutationName | null {
  if (!Array.isArray(mutationKey)) return null;
  const first = mutationKey[0];
  if (Array.isArray(first) && typeof first[0] === "string") {
    return first[0] as MutationName;
  }
  return null;
}

/**
 * Creates a MutationCache with a global onSuccess handler that automatically
 * invalidates the relevant queries for each tRPC mutation.
 */
export function createMutationCache(getQueryClient: () => QueryClient): MutationCache {
  return new MutationCache({
    onSuccess: (_data, _variables, _context, mutation) => {
      const name = extractProcedureName(mutation.options.mutationKey);
      if (!name) return;

      const queriesToInvalidate = INVALIDATION_MAP[name];
      if (!queriesToInvalidate) return;

      const queryClient = getQueryClient();
      for (const queryName of queriesToInvalidate) {
        void queryClient.invalidateQueries({
          queryKey: [[queryName]],
        });
      }
    },
  });
}