hp-prod-tracker/src/hooks/use-chat.ts

"use client";

import { useState, useCallback, useRef } from "react";
import { useQueryClient } from "@tanstack/react-query";
import { apiUrl } from "@/lib/api-client";

export interface ToolStatus {
  toolName: string;
  toolCallId: string;
  status: "running" | "done" | "error" | "pending_confirmation";
  input?: Record<string, any>;
  error?: string;
}

export interface PendingMutation {
  toolName: string;
  toolCallId: string;
  input: Record<string, any>;
  assistantMessage: string | null;
  messageId: string;
}

export interface EntitySuggestion {
  label: string;
  id: string;
  type: "project" | "deliverable" | "user" | "stage";
  description?: string;
}

export interface EntityCard {
  id: string;
  type: "project" | "deliverable" | "user" | "stage";
  name: string;
  status: string;
  priority?: string;
  code?: string;
  projectName?: string;
  projectId?: string;
  link: string;
  relevanceScore?: number;
}

export interface ChatMessage {
  id: string;
  role: "user" | "assistant";
  content: string;
  timestamp: Date;
  provider?: "claude";
  isLoading?: boolean;
  /** Tool calls in progress or completed during this assistant turn */
  toolCalls?: ToolStatus[];
  /** Clickable entity suggestions for disambiguation */
  suggestions?: EntitySuggestion[];
  /** Navigable entity cards extracted from tool results */
  entities?: EntityCard[];
}

export interface ChatContext {
  activeProjectId?: string;
  activeProjectName?: string;
  activeDeliverableId?: string;
  activeDeliverableName?: string;
}

/**
 * Parse an SSE stream and yield parsed events.
 */
async function* parseSSE(
  reader: ReadableStreamDefaultReader<Uint8Array>,
  decoder: TextDecoder
): AsyncGenerator<{ event: string; data: any }> {
  let buffer = "";

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;

    buffer += decoder.decode(value, { stream: true });
    const parts = buffer.split("\n\n");
    // Keep the last (possibly incomplete) chunk in the buffer
    buffer = parts.pop() || "";

    for (const part of parts) {
      if (!part.trim()) continue;

      let event = "message";
      let data = "";

      for (const line of part.split("\n")) {
        if (line.startsWith("event: ")) {
          event = line.slice(7);
        } else if (line.startsWith("data: ")) {
          data = line.slice(6);
        }
      }

      if (data) {
        try {
          yield { event, data: JSON.parse(data) };
        } catch {
          // Skip malformed JSON
        }
      }
    }
  }
}

export function useChat(context?: ChatContext) {
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  const [isLoading, setIsLoading] = useState(false);
  const [provider, setProvider] = useState<"claude" | "none">("none");
  const [pendingMutation, setPendingMutation] = useState<PendingMutation | null>(null);
  const queryClient = useQueryClient();
  const abortRef = useRef<AbortController | null>(null);
  // Keep a ref to messages for confirm/cancel callbacks
  const messagesRef = useRef(messages);
  messagesRef.current = messages;

  const sendMessage = useCallback(
    async (content: string) => {
      if (!content.trim() || isLoading) return;

      // Add user message
      const userMsg: ChatMessage = {
        id: `user-${Date.now()}`,
        role: "user",
        content: content.trim(),
        timestamp: new Date(),
      };

      // Add placeholder for assistant response
      const loadingMsg: ChatMessage = {
        id: `assistant-${Date.now()}`,
        role: "assistant",
        content: "",
        timestamp: new Date(),
        isLoading: true,
        toolCalls: [],
      };

      setMessages((prev) => [...prev, userMsg, loadingMsg]);
      setIsLoading(true);

      try {
        // Build message history for the API — send plain text only.
        // Tool interactions are resolved server-side within a single request,
        // so we never send tool_use/tool_result blocks back in history.
        // Trim long assistant responses to keep token count manageable.
        const apiMessages = [...messages, userMsg].map((m) => ({
          role: m.role,
          content: typeof m.content === "string" && m.content.length > 2000
            ? m.content.slice(0, 2000) + "..."
            : m.content,
        }));

        abortRef.current = new AbortController();

        const response = await fetch(apiUrl("/api/chat"), {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ messages: apiMessages, context }),
          signal: abortRef.current.signal,
        });

        if (!response.ok) {
          const errorData = await response.json().catch(() => null);
          throw new Error(
            errorData?.error || `Chat request failed: ${response.status}`
          );
        }

        // Parse SSE stream
        const reader = response.body!.getReader();
        const decoder = new TextDecoder();

        for await (const { event, data } of parseSSE(reader, decoder)) {
          switch (event) {
            case "tool_start":
              setMessages((prev) =>
                prev.map((m) =>
                  m.id === loadingMsg.id
                    ? {
                        ...m,
                        toolCalls: [
                          ...(m.toolCalls || []),
                          {
                            toolName: data.toolName,
                            toolCallId: data.toolCallId,
                            status: "running" as const,
                          },
                        ],
                      }
                    : m
                )
              );
              break;

            case "tool_end":
              setMessages((prev) =>
                prev.map((m) =>
                  m.id === loadingMsg.id
                    ? {
                        ...m,
                        toolCalls: (m.toolCalls || []).map((tc) =>
                          tc.toolCallId === data.toolCallId
                            ? {
                                ...tc,
                                status: data.success ? ("done" as const) : ("error" as const),
                                error: data.error,
                              }
                            : tc
                        ),
                      }
                    : m
                )
              );
              break;

            case "mutation_pending":
              // Update the tool call status to pending_confirmation
              setMessages((prev) =>
                prev.map((m) =>
                  m.id === loadingMsg.id
                    ? {
                        ...m,
                        content: data.assistantMessage || "",
                        isLoading: false,
                        toolCalls: (m.toolCalls || []).map((tc) =>
                          tc.toolCallId === data.toolCallId
                            ? { ...tc, status: "pending_confirmation" as const, input: data.input }
                            : tc
                        ),
                      }
                    : m
                )
              );
              setPendingMutation({
                toolName: data.toolName,
                toolCallId: data.toolCallId,
                input: data.input,
                assistantMessage: data.assistantMessage,
                messageId: loadingMsg.id,
              });
              break;

            case "message": {
              setMessages((prev) =>
                prev.map((m) =>
                  m.id === loadingMsg.id
                    ? {
                        ...m,
                        content: data.content,
                        provider: data.provider,
                        isLoading: false,
                        suggestions: data.suggestions,
                        entities: data.entities,
                      }
                    : m
                )
              );

              if (data.provider) {
                setProvider(data.provider);
              }

              // Invalidate TanStack Query caches for any mutated data
              if (data.invalidateKeys && data.invalidateKeys.length > 0) {
                for (const key of data.invalidateKeys) {
                  queryClient.invalidateQueries({ queryKey: [key] });
                }
                queryClient.invalidateQueries({ queryKey: ["dashboard"] });
              }
              break;
            }

            case "error":
              setMessages((prev) =>
                prev.map((m) =>
                  m.id === loadingMsg.id
                    ? {
                        ...m,
                        content: `Sorry, I encountered an error: ${data.error}`,
                        isLoading: false,
                      }
                    : m
                )
              );
              break;
          }
        }
      } catch (error) {
        if ((error as Error).name === "AbortError") return;

        setMessages((prev) =>
          prev.map((m) =>
            m.id === loadingMsg.id
              ? {
                  ...m,
                  content: `Sorry, I encountered an error: ${(error as Error).message}`,
                  isLoading: false,
                }
              : m
          )
        );
      } finally {
        setIsLoading(false);
        abortRef.current = null;
      }
    },
    [messages, isLoading, queryClient, context]
  );

  const clearMessages = useCallback(() => {
    setMessages([]);
  }, []);

  const cancelRequest = useCallback(() => {
    abortRef.current?.abort();
    setIsLoading(false);
    setMessages((prev) => prev.filter((m) => !m.isLoading));
  }, []);

  const confirmMutation = useCallback(() => {
    if (!pendingMutation) return;
    const pm = pendingMutation;
    setPendingMutation(null);

    // Update the tool status to "running" while we re-execute
    setMessages((prev) =>
      prev.map((m) =>
        m.id === pm.messageId
          ? {
              ...m,
              isLoading: true,
              toolCalls: (m.toolCalls || []).map((tc) =>
                tc.toolCallId === pm.toolCallId
                  ? { ...tc, status: "running" as const }
                  : tc
              ),
            }
          : m
      )
    );

    // Rebuild message history and add a confirmation marker
    // The server checks for [MUTATION_CONFIRMED] prefix to allow execution
    const apiMessages = messagesRef.current
      .filter((m) => m.id !== pm.messageId)
      .map((m) => ({ role: m.role, content: m.content }));

    apiMessages.push({
      role: "user",
      content: `[MUTATION_CONFIRMED] Execute ${pm.toolName}: ${JSON.stringify(pm.input)}`,
    });

    // Send as a new request
    (async () => {
      setIsLoading(true);
      try {
        abortRef.current = new AbortController();
        const response = await fetch(apiUrl("/api/chat"), {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ messages: apiMessages, context }),
          signal: abortRef.current.signal,
        });

        if (!response.ok) {
          const errorData = await response.json().catch(() => null);
          throw new Error(errorData?.error || `Request failed: ${response.status}`);
        }

        const reader = response.body!.getReader();
        const decoder = new TextDecoder();

        for await (const { event, data } of parseSSE(reader, decoder)) {
          switch (event) {
            case "tool_start":
              // Tool re-executing — already showing "running" status
              break;

            case "tool_end":
              setMessages((prev) =>
                prev.map((m) =>
                  m.id === pm.messageId
                    ? {
                        ...m,
                        toolCalls: (m.toolCalls || []).map((tc) =>
                          tc.toolCallId === data.toolCallId
                            ? {
                                ...tc,
                                status: data.success ? ("done" as const) : ("error" as const),
                                error: data.error,
                              }
                            : tc
                        ),
                      }
                    : m
                )
              );
              break;

            case "message":
              setMessages((prev) =>
                prev.map((m) =>
                  m.id === pm.messageId
                    ? {
                        ...m,
                        content: data.content,
                        provider: data.provider,
                        isLoading: false,
                        entities: data.entities,
                      }
                    : m
                )
              );
              if (data.invalidateKeys?.length > 0) {
                for (const key of data.invalidateKeys) {
                  queryClient.invalidateQueries({ queryKey: [key] });
                }
                queryClient.invalidateQueries({ queryKey: ["dashboard"] });
              }
              break;

            case "error":
              setMessages((prev) =>
                prev.map((m) =>
                  m.id === pm.messageId
                    ? {
                        ...m,
                        content: `Error: ${data.error}`,
                        isLoading: false,
                        toolCalls: (m.toolCalls || []).map((tc) =>
                          tc.toolCallId === pm.toolCallId
                            ? { ...tc, status: "error" as const, error: data.error }
                            : tc
                        ),
                      }
                    : m
                )
              );
              break;
          }
        }
      } catch (error) {
        if ((error as Error).name === "AbortError") return;
        setMessages((prev) =>
          prev.map((m) =>
            m.id === pm.messageId
              ? {
                  ...m,
                  content: `Error: ${(error as Error).message}`,
                  isLoading: false,
                }
              : m
          )
        );
      } finally {
        setIsLoading(false);
        abortRef.current = null;
      }
    })();
  }, [pendingMutation, queryClient, context]);

  const cancelMutation = useCallback(() => {
    if (!pendingMutation) return;
    const pm = pendingMutation;
    setPendingMutation(null);

    // Mark the tool as cancelled
    setMessages((prev) =>
      prev.map((m) =>
        m.id === pm.messageId
          ? {
              ...m,
              isLoading: false,
              content: (m.content || "") + "\n\n*Action cancelled by user.*",
              toolCalls: (m.toolCalls || []).map((tc) =>
                tc.toolCallId === pm.toolCallId
                  ? { ...tc, status: "error" as const, error: "Cancelled" }
                  : tc
              ),
            }
          : m
      )
    );
  }, [pendingMutation]);

  return {
    messages,
    isLoading,
    provider,
    pendingMutation,
    sendMessage,
    clearMessages,
    cancelRequest,
    confirmMutation,
    cancelMutation,
  };
}