FabledScribe/frontend/src/stores/chat.ts

import { ref, computed } from "vue";
import { defineStore } from "pinia";
import {
  apiGet,
  apiPost,
  apiPatch,
  apiDelete,
  apiSSEStream,
} from "@/api/client";
import { useToastStore } from "@/stores/toast";
import type {
  Conversation,
  ConversationDetail,
  ContextMeta,
  GenerationTiming,
  Message,
  OllamaStatus,
  SendMessageResponse,
  ToolCallRecord,
  ToolPendingRecord,
} from "@/types/chat";

interface ConvStreamState {
  streaming: boolean;
  content: string;
  thinking: string;
  toolCalls: ToolCallRecord[];
  status: string;
  pendingTool: ToolPendingRecord | null;
  contextMeta: ContextMeta | null;
}

interface QueuedMessage {
  content: string;
  contextNoteId?: number | null;
  includeNoteIds?: number[];
  think: boolean;
  contextNoteTitle?: string;
  excludeNoteIds?: number[];
  ragProjectId?: number | null;
  workspaceProjectId?: number | null;
}

export const useChatStore = defineStore("chat", () => {
  const conversations = ref<Conversation[]>([]);
  const currentConversation = ref<ConversationDetail | null>(null);
  const total = ref(0);
  const loading = ref(false);
  const convStreams = ref<Record<number, ConvStreamState>>({});
  const convQueues = ref<Record<number, QueuedMessage[]>>({});
  const ollamaStatus = ref<"checking" | "available" | "unavailable">("checking");
  const modelStatus = ref<"checking" | "loaded" | "cold" | "not_found">("checking");
  const defaultModel = ref("");
  let statusPollTimer: ReturnType<typeof setInterval> | null = null;

  function _getOrInitStream(id: number): ConvStreamState {
    if (!convStreams.value[id]) {
      convStreams.value[id] = {
        streaming: false, content: "", thinking: "", toolCalls: [],
        status: "", pendingTool: null, contextMeta: null,
      };
    }
    return convStreams.value[id];
  }

  // Per-conversation computed getters — same names as old refs, zero callers change
  const streaming = computed(() => {
    const id = currentConversation.value?.id;
    return !!id && (convStreams.value[id]?.streaming ?? false);
  });
  const streamingContent = computed(() => convStreams.value[currentConversation.value?.id ?? 0]?.content ?? "");
  const streamingThinking = computed(() => convStreams.value[currentConversation.value?.id ?? 0]?.thinking ?? "");
  const streamingToolCalls = computed(() => convStreams.value[currentConversation.value?.id ?? 0]?.toolCalls ?? []);
  const streamingStatus = computed(() => convStreams.value[currentConversation.value?.id ?? 0]?.status ?? "");
  const streamingPendingTool = computed(() => convStreams.value[currentConversation.value?.id ?? 0]?.pendingTool ?? null);
  const lastContextMeta = computed(() => convStreams.value[currentConversation.value?.id ?? 0]?.contextMeta ?? null);

  function isStreamingConv(id: number): boolean {
    return convStreams.value[id]?.streaming ?? false;
  }

  const queuedCount = computed(() => {
    const id = currentConversation.value?.id;
    return id ? (convQueues.value[id]?.length ?? 0) : 0;
  });

  const queuedMessages = computed(() => {
    const id = currentConversation.value?.id;
    return id ? (convQueues.value[id] ?? []) : [];
  });

  function clearQueue() {
    const id = currentConversation.value?.id;
    if (id) convQueues.value[id] = [];
  }

  // Chat is usable when Ollama is up and model is at least installed (cold starts still work)
  const chatReady = computed(
    () => ollamaStatus.value === "available" && (modelStatus.value === "loaded" || modelStatus.value === "cold")
  );

  async function fetchConversations(limit = 50, offset = 0) {
    loading.value = true;
    try {
      const data = await apiGet<{ conversations: Conversation[]; total: number }>(
        `/api/chat/conversations?limit=${limit}&offset=${offset}`
      );
      conversations.value = data.conversations;
      total.value = data.total;
    } catch (e) {
      useToastStore().show("Failed to load conversations", "error");
      throw e;
    } finally {
      loading.value = false;
    }
  }

  async function createConversation(title = ""): Promise<Conversation> {
    try {
      const conv = await apiPost<Conversation>("/api/chat/conversations", {
        title,
      });
      conversations.value.unshift(conv);
      return conv;
    } catch (e) {
      useToastStore().show("Failed to create conversation", "error");
      throw e;
    }
  }

  async function fetchConversation(id: number) {
    loading.value = true;
    try {
      currentConversation.value = await apiGet<ConversationDetail>(
        `/api/chat/conversations/${id}`
      );
    } catch (e) {
      useToastStore().show("Failed to load conversation", "error");
      throw e;
    } finally {
      loading.value = false;
    }
  }

  async function bulkDeleteConversations(ids: number[]) {
    if (!ids.length) return;
    await apiPost("/api/chat/conversations/bulk-delete", { ids });
    const idSet = new Set(ids);
    conversations.value = conversations.value.filter((c) => !idSet.has(c.id));
    if (currentConversation.value && idSet.has(currentConversation.value.id)) {
      currentConversation.value = null;
    }
    for (const id of ids) {
      delete convStreams.value[id];
      delete convQueues.value[id];
    }
  }

  async function deleteConversation(id: number) {
    try {
      await apiDelete(`/api/chat/conversations/${id}`);
      conversations.value = conversations.value.filter((c) => c.id !== id);
      if (currentConversation.value?.id === id) {
        currentConversation.value = null;
      }
      delete convStreams.value[id];
      delete convQueues.value[id];
    } catch (e) {
      useToastStore().show("Failed to delete conversation", "error");
      throw e;
    }
  }

  async function updateTitle(id: number, title: string) {
    try {
      const updated = await apiPatch<Conversation>(
        `/api/chat/conversations/${id}`,
        { title }
      );
      const idx = conversations.value.findIndex((c) => c.id === id);
      if (idx !== -1) {
        conversations.value[idx] = { ...conversations.value[idx], ...updated };
      }
      if (currentConversation.value?.id === id) {
        currentConversation.value.title = updated.title;
      }
    } catch (e) {
      useToastStore().show("Failed to update title", "error");
      throw e;
    }
  }

  async function sendMessage(
    content: string,
    contextNoteId?: number | null,
    includeNoteIds?: number[],
    think = false,
    contextNoteTitle?: string,
    excludeNoteIds?: number[],
    ragProjectId?: number | null,
    workspaceProjectId?: number | null,
  ) {
    if (!currentConversation.value) return;

    const convId = currentConversation.value.id;
    const s = _getOrInitStream(convId);

    // If already streaming, queue the message for after the current stream ends.
    if (s.streaming) {
      if (!convQueues.value[convId]) convQueues.value[convId] = [];
      convQueues.value[convId].push({
        content, contextNoteId, includeNoteIds, think, contextNoteTitle,
        excludeNoteIds, ragProjectId, workspaceProjectId,
      });
      return;
    }

    s.contextMeta = null;

    // If a write tool is waiting for confirmation, intercept single-word yes/no
    // responses rather than sending them as a new message.
    if (s.pendingTool) {
      const lower = content.trim().toLowerCase();
      const YES = new Set(["yes", "y", "ok", "sure", "proceed", "accept", "confirm", "do it", "go ahead", "yeah", "yep"]);
      const NO = new Set(["no", "n", "cancel", "decline", "stop", "nope", "abort", "don't"]);
      if (YES.has(lower)) {
        await confirmTool(true);
        return;
      }
      if (NO.has(lower)) {
        await confirmTool(false);
        return;
      }
    }

    // Add optimistic user message
    const userMsg: Message = {
      id: -Date.now(), // Temporary ID
      conversation_id: convId,
      role: "user",
      content,
      context_note_id: contextNoteId ?? null,
      context_note_title: contextNoteTitle ?? null,
      created_at: new Date().toISOString(),
    };
    currentConversation.value.messages.push(userMsg);

    s.streaming = true;
    s.content = "";

    // Phase 1: POST to start generation
    let assistantMessageId: number;
    try {
      const resp = await apiPost<SendMessageResponse>(
        `/api/chat/conversations/${convId}/messages`,
        {
          content,
          context_note_id: contextNoteId,
          include_note_ids: includeNoteIds?.length ? includeNoteIds : undefined,
          excluded_note_ids: excludeNoteIds?.length ? excludeNoteIds : undefined,
          think,
          rag_project_id: ragProjectId ?? undefined,
          workspace_project_id: workspaceProjectId ?? undefined,
        },
      );
      assistantMessageId = resp.assistant_message_id;
    } catch (e) {
      s.streaming = false;
      s.content = "";
      useToastStore().show("Failed to send message", "error");
      return;
    }

    // Immediately commit the message count so cleanup watcher never sees it as empty
    const convInList = conversations.value.find((c) => c.id === convId);
    if (convInList) convInList.message_count += 2;

    // Phase 2: Connect to SSE stream with reconnection
    await _streamGeneration(convId, assistantMessageId);
  }

  async function _streamGeneration(
    convId: number,
    assistantMessageId: number,
    initialLastEventId = -1,
    attempt = 0,
  ) {
    const MAX_RETRIES = 5;
    let lastEventId = initialLastEventId;
    let gotDone = false;

    const handle = apiSSEStream(
      `/api/chat/conversations/${convId}/generation/stream` +
        (lastEventId >= 0 ? `?last_event_id=${lastEventId}` : ""),
      (event) => {
        lastEventId = event.id;
        const s = convStreams.value[convId];
        if (!s) return;

        switch (event.event) {
          case "context":
            s.contextMeta = event.data.context as ContextMeta;
            break;
          case "status":
            s.status = event.data.status as string;
            break;
          case "thinking_chunk":
            s.thinking += event.data.chunk as string;
            break;
          case "chunk":
            s.content += event.data.chunk as string;
            s.status = "";
            break;
          case "tool_pending":
            s.pendingTool = event.data.tool_pending as ToolPendingRecord;
            break;
          case "tool_call":
            // Receiving a tool_call clears the pending confirmation card
            s.pendingTool = null;
            s.toolCalls = [...s.toolCalls, event.data.tool_call as ToolCallRecord];
            break;
          case "done":
            gotDone = true;
            {
              const assistantMsg: Message = {
                id: event.data.message_id as number,
                conversation_id: convId,
                role: "assistant",
                content: s.content,
                context_note_id: null,
                tool_calls: s.toolCalls.length ? [...s.toolCalls] : null,
                created_at: new Date().toISOString(),
                timing: event.data.timing as GenerationTiming | undefined,
                thinking: s.thinking || undefined,
              };
              if (currentConversation.value?.id === convId) {
                currentConversation.value.messages.push(assistantMsg);
              }
              // Update updated_at only — message_count was already incremented at send time
              const idx = conversations.value.findIndex((c) => c.id === convId);
              if (idx !== -1) {
                conversations.value[idx].updated_at = new Date().toISOString();
              }
              // Clear stream state
              s.content = "";
              s.thinking = "";
              s.toolCalls = [];
              s.status = "";
              s.pendingTool = null;
              s.streaming = false;
            }
            break;
          case "error":
            gotDone = true;
            s.streaming = false;
            s.content = "";
            s.thinking = "";
            s.toolCalls = [];
            s.status = "";
            s.pendingTool = null;
            if (currentConversation.value?.id === convId) {
              useToastStore().show(
                "Chat error: " + (event.data.error as string),
                "error",
              );
            }
            break;
        }
      },
    );

    // Wait for the stream to close (events are processed via callback above)
    await handle.done;

    // If stream ended without done/error, attempt reconnection regardless of which
    // conv is currently viewed — background streams retry independently
    if (!gotDone && attempt < MAX_RETRIES) {
      const delay = Math.min(1000 * 2 ** attempt, 10_000);
      await new Promise((r) => setTimeout(r, delay));
      return _streamGeneration(convId, assistantMessageId, lastEventId, attempt + 1);
    }

    // Recovery fallback: re-fetch conversation from DB only if currently viewing it
    if (!gotDone && currentConversation.value?.id === convId) {
      useToastStore().show(
        "Connection lost — response may be incomplete",
        "error",
      );
      try {
        await fetchConversation(convId);
      } catch {
        // Re-fetch failed — partial content visible on next page load
      }
    }

    // Final cleanup (guards against done/error not having fired)
    const s = convStreams.value[convId];
    if (s) {
      s.streaming = false;
      s.content = "";
      s.thinking = "";
      s.toolCalls = [];
      s.status = "";
      s.pendingTool = null;
    }

    // Process next queued message, if any.
    // Use setTimeout so this frame resolves before the next send begins.
    const queue = convQueues.value[convId];
    if (queue?.length && currentConversation.value?.id === convId) {
      const next = queue.shift()!;
      setTimeout(() => sendMessage(
        next.content,
        next.contextNoteId,
        next.includeNoteIds,
        next.think,
        next.contextNoteTitle,
        next.excludeNoteIds,
        next.ragProjectId,
        next.workspaceProjectId,
      ), 0);
    }
  }

  async function reconnectIfGenerating(convId: number): Promise<void> {
    // Skip if a stream is already active for this conversation
    if (isStreamingConv(convId)) return;

    const conv = currentConversation.value;
    if (!conv || conv.id !== convId) return;

    // Find the last assistant message still in generating state
    const lastMsg = [...conv.messages].reverse().find(
      (m) => m.role === "assistant" && m.status === "generating"
    );
    if (!lastMsg) return;

    // Remove the partial message — the done event will re-add the final version
    conv.messages = conv.messages.filter((m) => m.id !== lastMsg.id);

    const s = _getOrInitStream(convId);
    s.streaming = true;
    s.content = "";
    s.thinking = "";
    s.toolCalls = [];
    s.status = "Reconnecting...";
    s.pendingTool = null;
    s.contextMeta = null;

    // Reconnect from the beginning — the buffer replays all buffered events.
    // If the buffer is gone (>60s stale), _streamGeneration will exhaust retries
    // then fall back to fetchConversation to show the final saved content.
    await _streamGeneration(convId, lastMsg.id);
  }

  async function confirmTool(confirmed: boolean) {
    const convId = currentConversation.value?.id;
    if (!convId) return;
    // Optimistically clear so buttons disappear immediately
    const s = convStreams.value[convId];
    if (s) s.pendingTool = null;
    try {
      await apiPost(`/api/chat/conversations/${convId}/generation/confirm`, { confirmed });
    } catch {
      // Server will auto-decline on timeout — no action needed
    }
  }

  async function cancelGeneration() {
    if (!currentConversation.value) return;
    try {
      await apiPost(
        `/api/chat/conversations/${currentConversation.value.id}/generation/cancel`,
        {}
      );
    } catch {
      // Generation may have already finished — ignore
    }
  }

  async function saveMessageAsNote(messageId: number) {
    return await apiPost<Record<string, unknown>>(
      `/api/chat/messages/${messageId}/save-as-note`,
      {}
    );
  }

  async function summarizeAsNote(conversationId: number) {
    return await apiPost<Record<string, unknown>>(
      `/api/chat/conversations/${conversationId}/summarize`,
      {}
    );
  }

  async function fetchStatus() {
    try {
      const data = await apiGet<OllamaStatus>("/api/chat/status");
      ollamaStatus.value = data.ollama;
      modelStatus.value = data.model;
      defaultModel.value = data.default_model;
    } catch {
      ollamaStatus.value = "unavailable";
      modelStatus.value = "not_found";
    }
  }

  function startStatusPolling() {
    fetchStatus();
    stopStatusPolling();
    statusPollTimer = setInterval(fetchStatus, 30_000);
  }

  function stopStatusPolling() {
    if (statusPollTimer !== null) {
      clearInterval(statusPollTimer);
      statusPollTimer = null;
    }
  }

  async function warmModel(model: string) {
    try {
      await apiPost("/api/chat/warm", { model });
      // Poll faster after a warm request so the indicator updates promptly
      _pollUntilLoaded();
    } catch {
      // Warming is best-effort
    }
  }

  function _pollUntilLoaded() {
    let attempts = 0;
    const MAX = 12; // up to ~60s of fast polling
    const timer = setInterval(async () => {
      try {
        await fetchStatus();
      } catch {
        // best-effort polling
      }
      attempts++;
      if (modelStatus.value === "loaded" || attempts >= MAX) {
        clearInterval(timer);
      }
    }, 5_000);
  }

  return {
    conversations,
    currentConversation,
    total,
    loading,
    streaming,
    streamingContent,
    streamingThinking,
    streamingToolCalls,
    streamingStatus,
    streamingPendingTool,
    lastContextMeta,
    ollamaStatus,
    modelStatus,
    defaultModel,
    chatReady,
    isStreamingConv,
    queuedCount,
    queuedMessages,
    clearQueue,
    fetchConversations,
    createConversation,
    fetchConversation,
    deleteConversation,
    bulkDeleteConversations,
    updateTitle,
    sendMessage,
    reconnectIfGenerating,
    confirmTool,
    cancelGeneration,
    saveMessageAsNote,
    summarizeAsNote,
    warmModel,
    fetchStatus,
    startStatusPolling,
    stopStatusPolling,
  };
});