tidaldb/applications/iknowyou/lib/store.ts

import { create } from "zustand";
import { persist } from "zustand/middleware";
import type { ChatState, Conversation } from "./types";

function genId(): string {
  return crypto.randomUUID();
}

export const useChatStore = create<ChatState>()(
  persist(
    (set, get) => ({
      messages: [],
      isStreaming: false,
      error: null,

      conversations: [],
      activeConversationId: null,
      personId: crypto.randomUUID(),

      // --- Messages ---

      addUserMessage: (content: string) => {
        const id = genId();
        const now = Date.now();
        set((s) => {
          const updated = {
            messages: [
              ...s.messages,
              { id, role: "user" as const, content, timestamp: now },
            ],
            error: null,
          };

          // Update conversation title from first user message
          if (s.activeConversationId) {
            const conv = s.conversations.find(
              (c) => c.id === s.activeConversationId
            );
            if (conv && conv.title === "New conversation") {
              return {
                ...updated,
                conversations: s.conversations.map((c) =>
                  c.id === s.activeConversationId
                    ? {
                        ...c,
                        title: content.slice(0, 60),
                        lastMessageAt: now,
                      }
                    : c
                ),
              };
            }
            return {
              ...updated,
              conversations: s.conversations.map((c) =>
                c.id === s.activeConversationId
                  ? { ...c, lastMessageAt: now }
                  : c
              ),
            };
          }

          return updated;
        });
        return id;
      },

      startStreaming: () => {
        const id = genId();
        set((s) => ({
          messages: [
            ...s.messages,
            {
              id,
              role: "assistant" as const,
              content: "",
              timestamp: Date.now(),
            },
          ],
          isStreaming: true,
          error: null,
        }));
        return id;
      },

      appendToken: (id: string, token: string) => {
        set((s) => ({
          messages: s.messages.map((m) =>
            m.id === id ? { ...m, content: m.content + token } : m
          ),
        }));
      },

      finishStreaming: (id: string) => {
        set((s) => ({
          messages: s.messages.map((m) =>
            m.id === id ? { ...m, timestamp: Date.now() } : m
          ),
          isStreaming: false,
        }));
      },

      setError: (error: string | null) => {
        set({ error, isStreaming: false });
      },

      clearMessages: () => {
        set({ messages: [], error: null, isStreaming: false });
      },

      // --- Conversations ---

      createConversation: () => {
        const id = genId();
        const now = Date.now();
        const conv: Conversation = {
          id,
          title: "New conversation",
          createdAt: now,
          lastMessageAt: now,
        };
        set((s) => ({
          conversations: [conv, ...s.conversations],
          activeConversationId: id,
          messages: [],
          error: null,
          isStreaming: false,
        }));
        return id;
      },

      switchConversation: (id: string) => {
        const { activeConversationId } = get();
        if (id === activeConversationId) return;
        set({
          activeConversationId: id,
          messages: [],
          error: null,
          isStreaming: false,
        });
      },

      setMessages: (msgs) => {
        set({ messages: msgs });
      },

      updateConversationTitle: (id: string, title: string) => {
        set((s) => ({
          conversations: s.conversations.map((c) =>
            c.id === id ? { ...c, title: title.slice(0, 60) } : c
          ),
        }));
      },

      deleteConversation: (id: string) => {
        set((s) => {
          const filtered = s.conversations.filter((c) => c.id !== id);
          const wasActive = s.activeConversationId === id;
          return {
            conversations: filtered,
            activeConversationId: wasActive
              ? filtered[0]?.id ?? null
              : s.activeConversationId,
            messages: wasActive ? [] : s.messages,
          };
        });
      },

      switchPerson: () => {
        const newId = genId();
        set({
          personId: newId,
          conversations: [],
          activeConversationId: null,
          messages: [],
          error: null,
          isStreaming: false,
        });
        return newId;
      },
    }),
    {
      name: "aeries-conversations",
      partialize: (state) => ({
        conversations: state.conversations,
        activeConversationId: state.activeConversationId,
        personId: state.personId,
      }),
    }
  )
);