DreamChat/apps/frontend/src/stores/chatStore.ts

import { create } from 'zustand';
import {
  chatControllerCreateConversation,
  chatControllerGetConversations,
  chatControllerGetConversation,
  chatControllerDeleteConversation,
  chatControllerSendMessage,
} from '../api/generated/conversations/conversations';
import type { CreateConversationDto, SendMessageDto } from '../api/generated/model';
import type { Conversation, ConversationWithMessages, Message } from '../types';

interface ChatState {
  conversations: Conversation[];
  currentConversation: ConversationWithMessages | null;
  isLoading: boolean;
  isStreaming: boolean;
  error: string | null;

  // Actions
  fetchConversations: () => Promise<void>;
  getConversation: (id: string) => Promise<void>;
  createConversation: (data: CreateConversationDto) => Promise<Conversation>;
  deleteConversation: (id: string) => Promise<void>;
  sendMessage: (conversationId: string, data: SendMessageDto) => Promise<void>;
  addMessage: (message: Message) => void;
  setCurrentConversation: (conversation: ConversationWithMessages | null) => void;
  setStreaming: (isStreaming: boolean) => void;
  clearError: () => void;
}

export const useChatStore = create<ChatState>()((set) => ({
  conversations: [],
  currentConversation: null,
  isLoading: false,
  isStreaming: false,
  error: null,

  fetchConversations: async () => {
    set({ isLoading: true, error: null });
    try {
      const conversations = await chatControllerGetConversations() as unknown as Conversation[];
      set({ conversations, isLoading: false });
    } catch (error) {
      set({
        error: error instanceof Error ? error.message : 'Failed to fetch conversations',
        isLoading: false,
      });
    }
  },

  getConversation: async (id) => {
    set({ isLoading: true, error: null });
    try {
      const conversation = await chatControllerGetConversation(id) as unknown as ConversationWithMessages;
      set({ currentConversation: conversation, isLoading: false });
    } catch (error) {
      set({
        error: error instanceof Error ? error.message : 'Failed to fetch conversation',
        isLoading: false,
      });
    }
  },

  createConversation: async (data) => {
    set({ isLoading: true, error: null });
    try {
      const conversation = await chatControllerCreateConversation(data) as unknown as Conversation;
      set((state) => ({
        conversations: [conversation, ...state.conversations],
        isLoading: false,
      }));
      return conversation;
    } catch (error) {
      set({
        error: error instanceof Error ? error.message : 'Failed to create conversation',
        isLoading: false,
      });
      throw error;
    }
  },

  deleteConversation: async (id) => {
    set({ isLoading: true, error: null });
    try {
      await chatControllerDeleteConversation(id);
      set((state) => ({
        conversations: state.conversations.filter((c) => c.id !== id),
        currentConversation: state.currentConversation?.id === id ? null : state.currentConversation,
        isLoading: false,
      }));
    } catch (error) {
      set({
        error: error instanceof Error ? error.message : 'Failed to delete conversation',
        isLoading: false,
      });
    }
  },

  sendMessage: async (conversationId, data) => {
    set({ isLoading: true, error: null });
    try {
      const result = await chatControllerSendMessage(conversationId, data);
      
      set((state) => {
        if (!state.currentConversation) return state;
        
        const newMessages: Message[] = [
          result.userMessage as Message,
          result.assistantMessage as Message,
        ];
        
        return {
          currentConversation: {
            ...state.currentConversation,
            messages: [...state.currentConversation.messages, ...newMessages],
          },
          isLoading: false,
        };
      });
    } catch (error) {
      set({
        error: error instanceof Error ? error.message : 'Failed to send message',
        isLoading: false,
      });
    }
  },

  addMessage: (message) => {
    set((state) => {
      if (!state.currentConversation) return state;
      
      return {
        currentConversation: {
          ...state.currentConversation,
          messages: [...state.currentConversation.messages, message],
        },
      };
    });
  },

  setCurrentConversation: (conversation) => set({ currentConversation: conversation }),
  setStreaming: (isStreaming) => set({ isStreaming }),
  clearError: () => set({ error: null }),
}));