Codewalkers/packages/web/src/hooks/useRefineAgent.ts

import { useMemo, useCallback, useRef } from 'react';
import { trpc } from '@/lib/trpc';
import type { Agent, PendingQuestions } from '@codewalk-district/shared';

export type RefineAgentState = 'none' | 'running' | 'waiting' | 'completed' | 'crashed';

export interface ContentProposal {
  pageId: string;
  pageTitle: string;
  summary: string;
  markdown: string;
}

export interface SpawnRefineAgentOptions {
  initiativeId: string;
  instruction?: string;
}

export interface UseRefineAgentResult {
  /** Current refine agent for the initiative */
  agent: Agent | null;
  /** Current state of the refine agent */
  state: RefineAgentState;
  /** Questions from the agent (when state is 'waiting') */
  questions: PendingQuestions | null;
  /** Parsed content proposals (when state is 'completed') */
  proposals: ContentProposal[] | null;
  /** Raw result message (when state is 'completed') */
  result: string | null;
  /** Mutation for spawning a new refine agent */
  spawn: {
    mutate: (options: SpawnRefineAgentOptions) => void;
    isPending: boolean;
    error: Error | null;
  };
  /** Mutation for resuming agent with answers */
  resume: {
    mutate: (answers: Record<string, string>) => void;
    isPending: boolean;
    error: Error | null;
  };
  /** Whether any queries are loading */
  isLoading: boolean;
  /** Function to refresh agent data */
  refresh: () => void;
}

/**
 * Hook for managing refine agents for a specific initiative.
 *
 * Encapsulates the logic for finding, spawning, and interacting with refine agents
 * that analyze and suggest improvements to initiative content.
 *
 * @param initiativeId - The ID of the initiative to manage refine agents for
 * @returns Object with agent state, mutations, and helper functions
 *
 * @example
 * ```tsx
 * function RefineSection({ initiativeId }: { initiativeId: string }) {
 *   const {
 *     state,
 *     agent,
 *     questions,
 *     proposals,
 *     spawn,
 *     resume,
 *     refresh
 *   } = useRefineAgent(initiativeId);
 *
 *   const handleSpawn = () => {
 *     spawn.mutate({
 *       initiativeId,
 *       instruction: 'Focus on clarity and structure'
 *     });
 *   };
 *
 *   if (state === 'none') {
 *     return (
 *       <button onClick={handleSpawn} disabled={spawn.isPending}>
 *         Start Refine Agent
 *       </button>
 *     );
 *   }
 *
 *   if (state === 'waiting' && questions) {
 *     return (
 *       <QuestionForm
 *         questions={questions.questions}
 *         onSubmit={(answers) => resume.mutate(answers)}
 *         isSubmitting={resume.isPending}
 *       />
 *     );
 *   }
 *
 *   if (state === 'completed' && proposals) {
 *     return <ProposalReview proposals={proposals} onDismiss={refresh} />;
 *   }
 *
 *   return <div>Agent is {state}...</div>;
 * }
 * ```
 */
export function useRefineAgent(initiativeId: string): UseRefineAgentResult {
  const utils = trpc.useUtils();

  // Query all agents and find the active refine agent
  const agentsQuery = trpc.listAgents.useQuery();
  const agents = agentsQuery.data ?? [];

  const agent = useMemo(() => {
    // Find the most recent refine agent for this initiative
    const candidates = agents
      .filter(
        (a) =>
          a.mode === 'refine' &&
          a.initiativeId === initiativeId &&
          ['running', 'waiting_for_input', 'idle', 'crashed'].includes(a.status) &&
          !a.userDismissedAt, // Exclude dismissed agents
      )
      .sort(
        (a, b) =>
          new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime(),
      );
    return candidates[0] ?? null;
  }, [agents, initiativeId]);

  const state: RefineAgentState = useMemo(() => {
    if (!agent) return 'none';
    switch (agent.status) {
      case 'running':
        return 'running';
      case 'waiting_for_input':
        return 'waiting';
      case 'idle':
        return 'completed';
      case 'crashed':
        return 'crashed';
      default:
        return 'none';
    }
  }, [agent]);

  // Fetch questions when waiting for input
  const questionsQuery = trpc.getAgentQuestions.useQuery(
    { id: agent?.id ?? '' },
    { enabled: state === 'waiting' && !!agent },
  );

  // Fetch result when completed
  const resultQuery = trpc.getAgentResult.useQuery(
    { id: agent?.id ?? '' },
    { enabled: state === 'completed' && !!agent },
  );

  // Parse proposals from result
  const { proposals, result } = useMemo(() => {
    if (!resultQuery.data?.success || !resultQuery.data.message) {
      return { proposals: null, result: null };
    }

    const message = resultQuery.data.message;

    try {
      const parsed = JSON.parse(message);
      if (parsed.proposals && Array.isArray(parsed.proposals)) {
        const proposals: ContentProposal[] = parsed.proposals.map(
          (p: { pageId: string; title?: string; pageTitle?: string; summary: string; body?: string; markdown?: string }) => ({
            pageId: p.pageId,
            pageTitle: p.pageTitle ?? p.title ?? '',
            summary: p.summary,
            markdown: p.markdown ?? p.body ?? '',
          }),
        );
        return { proposals, result: message };
      }
    } catch {
      // Not JSON — treat as regular result
    }

    return { proposals: null, result: message };
  }, [resultQuery.data]);

  // Spawn mutation
  const spawnMutation = trpc.spawnArchitectRefine.useMutation({
    onSuccess: () => {
      void utils.listAgents.invalidate();
    },
  });

  // Resume mutation
  const resumeMutation = trpc.resumeAgent.useMutation({
    onSuccess: () => {
      void utils.listAgents.invalidate();
    },
  });

  // Keep mutation functions in refs so the returned spawn/resume objects are
  // stable across renders. tRPC mutation objects change identity every render,
  // which cascades into unstable callbacks → unstable props → Radix Dialog
  // re-renders that trigger the React 19 compose-refs infinite loop.
  const spawnMutateRef = useRef(spawnMutation.mutate);
  spawnMutateRef.current = spawnMutation.mutate;
  const agentRef = useRef(agent);
  agentRef.current = agent;
  const resumeMutateRef = useRef(resumeMutation.mutate);
  resumeMutateRef.current = resumeMutation.mutate;

  const spawnFn = useCallback(({ initiativeId, instruction }: SpawnRefineAgentOptions) => {
    spawnMutateRef.current({
      initiativeId,
      instruction: instruction?.trim() || undefined,
    });
  }, []);

  const spawn = useMemo(() => ({
    mutate: spawnFn,
    isPending: spawnMutation.isPending,
    error: spawnMutation.error,
  }), [spawnFn, spawnMutation.isPending, spawnMutation.error]);

  const resumeFn = useCallback((answers: Record<string, string>) => {
    const a = agentRef.current;
    if (a) {
      resumeMutateRef.current({ id: a.id, answers });
    }
  }, []);

  const resume = useMemo(() => ({
    mutate: resumeFn,
    isPending: resumeMutation.isPending,
    error: resumeMutation.error,
  }), [resumeFn, resumeMutation.isPending, resumeMutation.error]);

  const refresh = useCallback(() => {
    void utils.listAgents.invalidate();
  }, [utils]);

  const isLoading = agentsQuery.isLoading ||
    (state === 'waiting' && questionsQuery.isLoading) ||
    (state === 'completed' && resultQuery.isLoading);

  return {
    agent,
    state,
    questions: questionsQuery.data ?? null,
    proposals,
    result,
    spawn,
    resume,
    isLoading,
    refresh,
  };
}