Codewalkers/apps/web/src/hooks/useChatSession.ts

import { useCallback, useMemo, useRef, useState } from 'react';
import { toast } from 'sonner';
import { trpc } from '@/lib/trpc';
import { useLiveUpdates } from './useLiveUpdates';

export type ChatAgentStatus = 'idle' | 'running' | 'waiting' | 'crashed' | 'none';

export interface ChatMessage {
  id: string;
  chatSessionId: string;
  role: 'user' | 'assistant' | 'system';
  content: string;
  changeSetId: string | null;
  createdAt: string | Date;
}

export interface ChatSession {
  id: string;
  targetType: 'phase' | 'task';
  targetId: string;
  initiativeId: string;
  agentId: string | null;
  status: 'active' | 'closed';
  messages: ChatMessage[];
}

export interface UseChatSessionResult {
  session: ChatSession | null;
  messages: ChatMessage[];
  agentStatus: ChatAgentStatus;
  agentError: string | null;
  sendMessage: (message: string) => void;
  retryLastMessage: () => void;
  closeSession: () => void;
  startNewSession: () => void;
  isSending: boolean;
  isLoading: boolean;
}

/**
 * Hook for managing a chat session with a phase or task.
 *
 * Queries the active chat session, tracks agent status,
 * and provides mutations for sending messages and closing.
 */
export function useChatSession(
  targetType: 'phase' | 'task',
  targetId: string,
  initiativeId: string,
): UseChatSessionResult {
  const utils = trpc.useUtils();

  // Optimistic messages shown before server confirms
  const [optimisticMessages, setOptimisticMessages] = useState<ChatMessage[]>([]);

  // Live updates for chat + agent events
  useLiveUpdates([
    { prefix: 'chat:', invalidate: ['getChatSession'] },
    { prefix: 'agent:', invalidate: ['getChatSession', 'getAgent'] },
    { prefix: 'changeset:', invalidate: ['getChatSession'] },
  ]);

  // Query active session
  const sessionQuery = trpc.getChatSession.useQuery(
    { targetType, targetId },
    { enabled: !!targetId },
  );
  const session = (sessionQuery.data as ChatSession | null) ?? null;
  const serverMessages = session?.messages ?? [];

  // Merge: show server messages, plus any optimistic ones not yet confirmed
  const serverMsgIds = new Set(serverMessages.map(m => m.id));
  const messages = [
    ...serverMessages,
    ...optimisticMessages.filter(m => !serverMsgIds.has(m.id)),
  ];

  // Query agent status if session has an agent
  const agentQuery = trpc.getAgent.useQuery(
    { id: session?.agentId ?? '' },
    { enabled: !!session?.agentId },
  );

  const agentStatus: ChatAgentStatus = useMemo(() => {
    if (!session?.agentId || !agentQuery.data) return 'none';
    const status = agentQuery.data.status;
    switch (status) {
      case 'running':
        return 'running';
      case 'waiting_for_input':
        return 'waiting';
      case 'idle':
        return 'idle';
      case 'crashed':
        return 'crashed';
      default:
        return 'none';
    }
  }, [session?.agentId, agentQuery.data]);

  // Extract error message when agent crashed
  const agentError = useMemo(() => {
    if (agentStatus !== 'crashed' || !agentQuery.data) return null;
    const result = (agentQuery.data as { result?: string }).result;
    if (!result) return 'Agent crashed unexpectedly';
    try {
      const parsed = JSON.parse(result);
      return parsed.message ?? result;
    } catch {
      return result;
    }
  }, [agentStatus, agentQuery.data]);

  // Send message mutation
  const sendMutation = trpc.sendChatMessage.useMutation({
    onSuccess: () => {
      void utils.getChatSession.invalidate({ targetType, targetId });
      // Clear optimistic messages once server confirms
      setOptimisticMessages([]);
    },
    onError: (err) => {
      toast.error(`Chat failed: ${err.message}`);
      // Remove optimistic messages on error
      setOptimisticMessages([]);
    },
  });

  // Close session mutation
  const closeMutation = trpc.closeChatSession.useMutation({
    onSuccess: () => {
      void utils.getChatSession.invalidate({ targetType, targetId });
    },
    onError: (err) => {
      toast.error(`Failed to close chat: ${err.message}`);
    },
  });

  const sendMutateRef = useRef(sendMutation.mutate);
  sendMutateRef.current = sendMutation.mutate;
  const closeMutateRef = useRef(closeMutation.mutate);
  closeMutateRef.current = closeMutation.mutate;
  const sessionRef = useRef(session);
  sessionRef.current = session;

  const sendMessage = useCallback(
    (message: string) => {
      // Add optimistic user message immediately
      const optimistic: ChatMessage = {
        id: `opt-${Date.now()}`,
        chatSessionId: '',
        role: 'user',
        content: message,
        changeSetId: null,
        createdAt: new Date().toISOString(),
      };
      setOptimisticMessages(prev => [...prev, optimistic]);

      sendMutateRef.current({
        targetType,
        targetId,
        initiativeId,
        message,
      });
    },
    [targetType, targetId, initiativeId],
  );

  const messagesRef = useRef(messages);
  messagesRef.current = messages;

  const retryLastMessage = useCallback(() => {
    const lastUserMsg = [...messagesRef.current].reverse().find(m => m.role === 'user');
    if (!lastUserMsg) return;
    sendMutateRef.current({
      targetType,
      targetId,
      initiativeId,
      message: lastUserMsg.content,
      retry: true,
    });
  }, [targetType, targetId, initiativeId]);

  const closeSession = useCallback(() => {
    const s = sessionRef.current;
    if (s) {
      closeMutateRef.current({ sessionId: s.id });
    }
  }, []);

  const startNewSession = useCallback(() => {
    const s = sessionRef.current;
    if (!s) return;
    closeMutateRef.current({ sessionId: s.id });
    setOptimisticMessages([]);
  }, []);

  const isLoading = sessionQuery.isLoading;
  const isSending = sendMutation.isPending;

  return {
    session,
    messages,
    agentStatus,
    agentError,
    sendMessage,
    retryLastMessage,
    closeSession,
    startNewSession,
    isSending,
    isLoading,
  };
}