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,
  Message,
  OllamaStatus,
  SendMessageResponse,
  ToolCallRecord,
} from "@/types/chat";

export const useChatStore = defineStore("chat", () => {
  const conversations = ref<Conversation[]>([]);
  const currentConversation = ref<ConversationDetail | null>(null);
  const total = ref(0);
  const loading = ref(false);
  const streaming = ref(false);
  const streamingContent = ref("");
  const streamingToolCalls = ref<ToolCallRecord[]>([]);
  const lastContextMeta = ref<ContextMeta | null>(null);
  const ollamaStatus = ref<"checking" | "available" | "unavailable">("checking");
  const modelStatus = ref<"checking" | "ready" | "not_found">("checking");
  const defaultModel = ref("");
  let statusPollTimer: ReturnType<typeof setInterval> | null = null;

  const chatReady = computed(
    () => ollamaStatus.value === "available" && modelStatus.value === "ready"
  );

  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 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;
      }
    } 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,
    excludeNoteIds?: number[],
  ) {
    if (!currentConversation.value) return;

    const convId = currentConversation.value.id;
    lastContextMeta.value = null;

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

    streaming.value = true;
    streamingContent.value = "";

    // 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,
          exclude_note_ids: excludeNoteIds?.length ? excludeNoteIds : undefined,
        },
      );
      assistantMessageId = resp.assistant_message_id;
    } catch (e) {
      streaming.value = false;
      streamingContent.value = "";
      useToastStore().show("Failed to send message", "error");
      return;
    }

    // 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;

        switch (event.event) {
          case "context":
            lastContextMeta.value = event.data.context as ContextMeta;
            break;
          case "chunk":
            streamingContent.value += event.data.chunk as string;
            break;
          case "tool_call":
            streamingToolCalls.value = [
              ...streamingToolCalls.value,
              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: streamingContent.value,
                context_note_id: null,
                tool_calls: streamingToolCalls.value.length
                  ? [...streamingToolCalls.value]
                  : null,
                created_at: new Date().toISOString(),
              };
              if (currentConversation.value?.id === convId) {
                currentConversation.value.messages.push(assistantMsg);
              }
              streamingContent.value = "";
              streamingToolCalls.value = [];
              streaming.value = false;

              // Update conversation in list
              const idx = conversations.value.findIndex((c) => c.id === convId);
              if (idx !== -1) {
                conversations.value[idx].message_count += 2;
                conversations.value[idx].updated_at = new Date().toISOString();
              }
            }
            break;
          case "error":
            gotDone = true;
            streaming.value = false;
            streamingContent.value = "";
            streamingToolCalls.value = [];
            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
    if (!gotDone && attempt < MAX_RETRIES) {
      const delay = Math.min(1000 * 2 ** attempt, 10_000);
      await new Promise((r) => setTimeout(r, delay));
      if (currentConversation.value?.id === convId) {
        return _streamGeneration(convId, assistantMessageId, lastEventId, attempt + 1);
      }
    }

    // Recovery fallback: re-fetch conversation from DB
    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
      }
    }

    streaming.value = false;
    streamingContent.value = "";
    streamingToolCalls.value = [];
  }

  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 });
    } catch {
      // Warming is best-effort
    }
  }

  return {
    conversations,
    currentConversation,
    total,
    loading,
    streaming,
    streamingContent,
    streamingToolCalls,
    lastContextMeta,
    ollamaStatus,
    modelStatus,
    defaultModel,
    chatReady,
    fetchConversations,
    createConversation,
    fetchConversation,
    deleteConversation,
    updateTitle,
    sendMessage,
    cancelGeneration,
    saveMessageAsNote,
    summarizeAsNote,
    warmModel,
    fetchStatus,
    startStatusPolling,
    stopStatusPolling,
  };
});