client/src/hooks/useChat.ts

import { useState, useCallback, useRef } from "react";

export interface ToolCallInfo {
  id: string;
  name: string;
  arguments: string;
  result?: string;
  isError?: boolean;
  durationMs?: number;
  isExecuting?: boolean;
}

/** A content segment preserving the order of text and tool calls as they arrive */
export type ContentSegment =
  | { type: "text"; content: string }
  | { type: "tool"; toolIndex: number }; // index into toolCalls[]

export interface ChatMessage {
  id?: number;
  role: "user" | "assistant" | "system" | "tool";
  content: string;
  toolCalls?: ToolCallInfo[];
  /** Ordered segments preserving interleaved text + tool call positions */
  contentSegments?: ContentSegment[];
  thinkingBlocks?: { thinking: string; durationMs?: number }[];
  isStreaming?: boolean;
  promptTokens?: number;
  completionTokens?: number;
  cost?: number;
  model?: string;
}

export interface UseChatReturn {
  messages: ChatMessage[];
  isStreaming: boolean;
  error: string | null;
  sendMessage: (content: string, sessionId: number) => void;
  stopStreaming: () => void;
  clearMessages: () => void;
  setMessages: React.Dispatch<React.SetStateAction<ChatMessage[]>>;
  currentStatus: string;
}

export function useChat(): UseChatReturn {
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  const [isStreaming, setIsStreaming] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [currentStatus, setCurrentStatus] = useState("");
  const abortRef = useRef<AbortController | null>(null);

  const stopStreaming = useCallback(() => {
    if (abortRef.current) {
      abortRef.current.abort();
      abortRef.current = null;
    }
    setIsStreaming(false);
    setCurrentStatus("");
  }, []);

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

  const sendMessage = useCallback(
    async (content: string, sessionId: number) => {
      if (isStreaming) return;

      setError(null);
      setIsStreaming(true);
      setCurrentStatus("Sending...");

      // Add user message
      const userMsg: ChatMessage = { role: "user", content };
      setMessages((prev) => [...prev, userMsg]);

      // Create assistant placeholder
      const assistantMsg: ChatMessage = {
        role: "assistant",
        content: "",
        toolCalls: [],
        contentSegments: [],
        isStreaming: true,
      };
      setMessages((prev) => [...prev, assistantMsg]);

      const controller = new AbortController();
      abortRef.current = controller;

      try {
        const response = await fetch("/api/chat/stream", {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ sessionId, message: content }),
          signal: controller.signal,
        });

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

        const reader = response.body?.getReader();
        if (!reader) throw new Error("No response body");

        const decoder = new TextDecoder();
        let buffer = "";
        let currentToolCalls: ToolCallInfo[] = [];
        let assistantContent = "";
        let lastEventType = "";

        // ─── Segment tracking ───────────────────────────────────────
        // Track ordered segments: text chunks and tool call references
        // so MessageBubble can render them interleaved.
        let segments: ContentSegment[] = [];
        // Track whether the last segment is a text segment (to append to it)
        let lastSegmentIsText = false;
        // Track text accumulated since last tool call (for current text segment)
        let currentTextSegmentStart = 0; // char index into assistantContent
        let pendingAssistantMessage = ""; // accumulates SendUserMessage/Brief text

        const updateMessage = () => {
          setMessages((prev) => {
            const updated = [...prev];
            const last = updated[updated.length - 1];
            if (last && last.role === "assistant") {
              updated[updated.length - 1] = {
                ...last,
                content: assistantContent,
                toolCalls: [...currentToolCalls],
                contentSegments: [...segments],
              };
            }
            return updated;
          });
        };

        // Helper: ensure current text is captured in segments
        const flushTextSegment = () => {
          const currentText = assistantContent.slice(currentTextSegmentStart).trim();
          if (currentText && lastSegmentIsText) {
            // Update the last text segment with full current text
            segments[segments.length - 1] = {
              type: "text",
              content: assistantContent.slice(currentTextSegmentStart),
            };
          }
        };

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

          buffer += decoder.decode(value, { stream: true });
          const lines = buffer.split("\n");
          buffer = lines.pop() || "";

          for (const line of lines) {
            // Track event type from SSE event: lines
            if (line.startsWith("event: ")) {
              lastEventType = line.slice(7).trim();
              continue;
            }
            if (!line.startsWith("data: ")) continue;

            const dataStr = line.slice(6).trim();
            if (!dataStr || dataStr === "{}") continue;

            try {
              const data = JSON.parse(dataStr);

              // Route by event type first, then fall back to data shape
              switch (lastEventType) {
                case "text_delta": {
                  const newText = data.text || "";
                  assistantContent += newText;

                  // If we don't have a text segment at the end, create one
                  if (!lastSegmentIsText) {
                    currentTextSegmentStart = assistantContent.length - newText.length;
                    segments.push({
                      type: "text",
                      content: newText,
                    });
                    lastSegmentIsText = true;
                  } else {
                    // Append to existing text segment
                    segments[segments.length - 1] = {
                      type: "text",
                      content: assistantContent.slice(currentTextSegmentStart),
                    };
                  }
                  updateMessage();
                  break;
                }

                case "tool_call_start": {
                  // Flush any pending text segment
                  flushTextSegment();

                  const toolIndex = currentToolCalls.length;
                  currentToolCalls.push({
                    id: data.id,
                    name: data.name,
                    arguments: "",
                    isExecuting: true,
                  });

                  // Add tool segment
                  segments.push({ type: "tool", toolIndex });
                  lastSegmentIsText = false;

                  updateMessage();
                  break;
                }

                case "tool_call_delta": {
                  const tcIdx = currentToolCalls.findIndex((t) => t.id === data.id);
                  if (tcIdx >= 0) {
                    currentToolCalls[tcIdx] = {
                      ...currentToolCalls[tcIdx],
                      arguments: (currentToolCalls[tcIdx].arguments || "") + (data.arguments || ""),
                    };
                  }
                  break;
                }

                case "tool_call_end":
                  // Tool call arguments finalized
                  break;

                case "tool_result": {
                  const tcr = currentToolCalls.find(
                    (t) => t.id === data.toolCallId
                  );
                  if (tcr) {
                    tcr.result = data.output;
                    tcr.isError = data.isError || false;
                    tcr.durationMs = data.durationMs || 0;
                    tcr.isExecuting = false;
                  }
                  updateMessage();
                  break;
                }

                case "message_start":
                  // Bug #7 fix: Do NOT reset assistantContent between iterations.
                  // Original claw-code accumulates text across tool call iterations.
                  // Resetting here loses text the LLM wrote before calling tools.
                  break;

                case "message_end":
                  setMessages((prev) => {
                    const updated = [...prev];
                    const last = updated[updated.length - 1];
                    if (last && last.role === "assistant") {
                      updated[updated.length - 1] = {
                        ...last,
                        isStreaming: false,
                        promptTokens: data.promptTokens,
                        completionTokens: data.completionTokens,
                        cost: data.cost,
                        model: data.model,
                      };
                    }
                    return updated;
                  });
                  break;

                case "thinking_delta": {
                  // Accumulate thinking blocks from extended_thinking
                  setMessages((prev) => {
                    const updated = [...prev];
                    const last = updated[updated.length - 1];
                    if (last && last.role === "assistant") {
                      const blocks = last.thinkingBlocks ? [...last.thinkingBlocks] : [];
                      if (blocks.length === 0 || data.isNew) {
                        blocks.push({ thinking: data.text || "", durationMs: data.durationMs });
                      } else {
                        blocks[blocks.length - 1] = {
                          ...blocks[blocks.length - 1],
                          thinking: blocks[blocks.length - 1].thinking + (data.text || ""),
                          durationMs: data.durationMs || blocks[blocks.length - 1].durationMs,
                        };
                      }
                      updated[updated.length - 1] = { ...last, thinkingBlocks: blocks };
                    }
                    return updated;
                  });
                  break;
                }

                case "status":
                  setCurrentStatus(data.message || data.status || "");
                  break;

                case "title_update":
                  window.dispatchEvent(
                    new CustomEvent("session-title-update", {
                      detail: { title: data.title },
                    })
                  );
                  break;

                case "ask_user":
                  // Dispatch event for Home.tsx to show ask dialog
                  window.dispatchEvent(
                    new CustomEvent("agent-ask-user", {
                      detail: { question: data.question },
                    })
                  );
                  break;

                case "assistant_message":
                  // Model sent a progress message via SendUserMessage/Brief.
                  // Display it as inline text in the current message.
                  if (data.message) {
                    const msgText = data.message;
                    assistantContent += (assistantContent ? "\n" : "") + msgText;

                    // Create a new text segment for this progress message
                    if (!lastSegmentIsText) {
                      currentTextSegmentStart = assistantContent.length - msgText.length;
                      segments.push({ type: "text", content: msgText });
                      lastSegmentIsText = true;
                    } else {
                      segments[segments.length - 1] = {
                        type: "text",
                        content: assistantContent.slice(currentTextSegmentStart),
                      };
                    }
                    updateMessage();
                  }
                  break;

                case "plan_update":
                  // Dispatch event for Home.tsx to update plan state
                  window.dispatchEvent(
                    new CustomEvent("agent-plan-update", {
                      detail: {
                        active: data.active,
                        steps: data.steps,
                      },
                    })
                  );
                  break;

                case "auto_compact":
                  // Auto-compaction notification — matches original format_compact_report
                  window.dispatchEvent(
                    new CustomEvent("auto-compact", {
                      detail: {
                        removedMessages: data.removedCount || data.removedMessages,
                        remainingMessages: data.keptCount || data.remainingMessages,
                        summary: data.summary,
                      },
                    })
                  );
                  // Add system message showing compact report
                  setMessages((prev) => [
                    ...prev,
                    {
                      role: "system" as const,
                      content: `Compact\n  Result           compacted\n  Messages removed ${data.removedCount || data.removedMessages}\n  Messages kept    ${data.keptCount || data.remainingMessages}`,
                    },
                  ]);
                  break;

                case "context_usage":
                  // Context window usage tracking — display in status bar
                  setCurrentStatus(
                    `Context: ${data.usagePercent}% (${data.estimatedTokens?.toLocaleString()}/${data.contextWindow?.toLocaleString()} tokens, ${data.messageCount} msgs)`
                  );
                  break;

                case "buddy_event":
                  // Buddy system events — dispatch to BuddySprite for XP/progression
                  window.dispatchEvent(
                    new CustomEvent("buddy-event", {
                      detail: {
                        type: data.type,
                        toolName: data.toolName,
                        iteration: data.iteration,
                        toolCallCount: data.toolCallCount,
                        lines: data.lines,
                      },
                    })
                  );
                  break;

                case "error":
                  setError(data.message || "Unknown error");
                  // Bug #6 fix: Finalize the assistant message so it doesn't look stuck.
                  setMessages((prev) => {
                    const updated = [...prev];
                    const last = updated[updated.length - 1];
                    if (last && last.role === "assistant") {
                      updated[updated.length - 1] = {
                        ...last,
                        isStreaming: false,
                        content: last.content || `[Error: ${data.message || "Unknown error"}]`,
                      };
                    }
                    return updated;
                  });
                  break;

                case "usage":
                  // Usage stats — update last message
                  setMessages((prev) => {
                    const updated = [...prev];
                    const last = updated[updated.length - 1];
                    if (last && last.role === "assistant") {
                      updated[updated.length - 1] = {
                        ...last,
                        promptTokens: data.promptTokens,
                        completionTokens: data.completionTokens,
                        cost: data.cost,
                      };
                    }
                    return updated;
                  });
                  break;

                default:
                  // Fallback: detect event type from data shape
                  handleSSEDataFallback(data, {
                    onTextDelta: (text: string) => {
                      assistantContent += text;
                      if (!lastSegmentIsText) {
                        currentTextSegmentStart = assistantContent.length - text.length;
                        segments.push({ type: "text", content: text });
                        lastSegmentIsText = true;
                      } else {
                        segments[segments.length - 1] = {
                          type: "text",
                          content: assistantContent.slice(currentTextSegmentStart),
                        };
                      }
                      updateMessage();
                    },
                    onToolCallStart: (id: string, name: string) => {
                      flushTextSegment();
                      const toolIndex = currentToolCalls.length;
                      currentToolCalls.push({
                        id,
                        name,
                        arguments: "",
                        isExecuting: true,
                      });
                      segments.push({ type: "tool", toolIndex });
                      lastSegmentIsText = false;
                      updateMessage();
                    },
                    onToolCallDelta: (id: string, args: string) => {
                      const tcIdx2 = currentToolCalls.findIndex((t) => t.id === id);
                      if (tcIdx2 >= 0) {
                        currentToolCalls[tcIdx2] = {
                          ...currentToolCalls[tcIdx2],
                          arguments: (currentToolCalls[tcIdx2].arguments || "") + args,
                        };
                      }
                    },
                    onToolResult: (
                      toolCallId: string,
                      _toolName: string,
                      output: string,
                      isErr: boolean,
                      durationMs: number
                    ) => {
                      const tc2 = currentToolCalls.find(
                        (t) => t.id === toolCallId
                      );
                      if (tc2) {
                        tc2.result = output;
                        tc2.isError = isErr;
                        tc2.durationMs = durationMs;
                        tc2.isExecuting = false;
                      }
                      updateMessage();
                    },
                    onMessageEnd: (d: any) => {
                      setMessages((prev) => {
                        const updated = [...prev];
                        const last = updated[updated.length - 1];
                        if (last && last.role === "assistant") {
                          updated[updated.length - 1] = {
                            ...last,
                            isStreaming: false,
                            promptTokens: d.promptTokens,
                            completionTokens: d.completionTokens,
                            cost: d.cost,
                            model: d.model,
                          };
                        }
                        return updated;
                      });
                    },
                    onStatus: (status: string, message?: string) => {
                      setCurrentStatus(message || status);
                    },
                    onTitleUpdate: (title: string) => {
                      window.dispatchEvent(
                        new CustomEvent("session-title-update", {
                          detail: { title },
                        })
                      );
                    },
                    onError: (msg: string) => {
                      setError(msg);
                    },
                  });
                  break;
              }

              // Reset event type after processing
              lastEventType = "";
            } catch {
              // Skip malformed JSON
            }
          }
        }
      } catch (err: any) {
        if (err.name !== "AbortError") {
          setError(err.message || "Connection failed");
          setMessages((prev) => {
            const last = prev[prev.length - 1];
            if (last?.role === "assistant" && !last.content) {
              return prev.slice(0, -1);
            }
            return prev;
          });
        }
      } finally {
        setIsStreaming(false);
        setCurrentStatus("");
        abortRef.current = null;
        setMessages((prev) => {
          const updated = [...prev];
          const last = updated[updated.length - 1];
          if (last && last.isStreaming) {
            updated[updated.length - 1] = { ...last, isStreaming: false };
          }
          return updated;
        });
      }
    },
    [isStreaming]
  );

  return {
    messages,
    isStreaming,
    error,
    sendMessage,
    stopStreaming,
    clearMessages,
    setMessages,
    currentStatus,
  };
}

// Fallback: parse SSE events based on data structure (when event type is missing)
function handleSSEDataFallback(
  data: any,
  handlers: {
    onTextDelta: (text: string) => void;
    onToolCallStart: (id: string, name: string) => void;
    onToolCallDelta: (id: string, args: string) => void;
    onToolResult: (
      toolCallId: string,
      toolName: string,
      output: string,
      isError: boolean,
      durationMs: number
    ) => void;
    onMessageEnd: (data: any) => void;
    onStatus: (status: string, message?: string) => void;
    onTitleUpdate: (title: string) => void;
    onError: (message: string) => void;
  }
) {
  if (data.text !== undefined) {
    handlers.onTextDelta(data.text);
  } else if (data.toolCallId !== undefined && data.output !== undefined) {
    handlers.onToolResult(
      data.toolCallId,
      data.toolName || "",
      data.output,
      data.isError || false,
      data.durationMs || 0
    );
  } else if (
    data.name !== undefined &&
    data.arguments !== undefined &&
    data.id !== undefined
  ) {
    if (data.arguments === "") {
      handlers.onToolCallStart(data.id, data.name);
    } else {
      handlers.onToolCallDelta(data.id, data.arguments);
    }
  } else if (
    data.id !== undefined &&
    data.name !== undefined &&
    !data.arguments
  ) {
    handlers.onToolCallStart(data.id, data.name);
  } else if (
    data.promptTokens !== undefined ||
    data.completionTokens !== undefined
  ) {
    handlers.onMessageEnd(data);
  } else if (data.status !== undefined) {
    handlers.onStatus(data.status, data.message);
  } else if (data.title !== undefined) {
    handlers.onTitleUpdate(data.title);
  } else if (data.message !== undefined && !data.status) {
    handlers.onError(data.message);
  }
}