neoban/frontend/src/app/chat/chat.store.ts

import { patchState, signalStore, withComputed, withMethods, withState } from '@ngrx/signals';
import { computed, inject } from '@angular/core';
import { HttpErrorResponse } from '@angular/common/http';
import { ChatService } from './chat-service';
import {
  MessageState,
  Message,
  MessageResponse,
  ChatResponse,
  ChatState,
  Chat,
} from './chat.types';
import { debounceTime, distinctUntilChanged, lastValueFrom, of, pipe, switchMap } from 'rxjs';
import { Router } from '@angular/router';
import { JsonApiResource } from '../core/types/api';
import { environment } from '../../environments/environment.development';
import { rxMethod } from '@ngrx/signals/rxjs-interop';

const initialMessageState: MessageState = {
  messages: [
    {
      id: 'null',
      type: 'messages',
      attributes: {
        role: 'assistant',
        content: "What's you want to post today?",
        attachments: [],
        createdAt: new Date(),
      },
    },
  ],
  isLoading: true,
  isAgentThinking: false,
  isLoadingPaginated: false,
  currentPage: 1,
  lastPage: 1,
  id: null,
};

const initalChatState: ChatState = {
  _chats: [],
  _serverSearchResults: [],
  _filteredChats: [],
  isLoading: true,
  searchQuery: null,
};

export const MessageStore = signalStore(
  { providedIn: 'root' },
  withState(initialMessageState),
  withMethods((store) => {
    const chatService = inject(ChatService);
    const router = inject(Router);

    const newChat = async () => {
      patchState(store, { isLoading: true });
      try {
        const response = await lastValueFrom(chatService.newChat());
        const chatId = response.data.id;

        patchState(store, { id: chatId, isLoading: false });
        await router.navigate(['/', chatId], { replaceUrl: true });
        return chatId;
      } catch (error) {
        patchState(store, { isLoading: false });
        console.error('Failed to create a new chat:', error);
        throw error;
      }
    };
    const setErrorMessage = (message: string, role: 'assistant' | 'user' = 'user') => {
      const errorMessage: JsonApiResource<Message> = {
        id: Date.now().toString(),
        type: 'messages',
        attributes: {
          role: role,
          content: message,
          attachments: [],
          createdAt: new Date(),
        },
      };
      patchState(store, (state) => ({
        messages: [...state.messages, errorMessage],
      }));
    };

    return {
      newChat,
      setChatId: (chatId: string) => {
        patchState(store, { id: chatId });
      },
      sendMessage: async (content: string) => {
        if (!store.id()) {
          try {
            await newChat();
          } catch (error) {
            // If chat creation fails, append an error message and stop execution
            setErrorMessage(
              'Failed to initialize a new chat session. Please try again.',
              'assistant',
            );
            return;
          }
        }

        // Add user message
        const userMessage: JsonApiResource<Message> = {
          id: Date.now().toString(),
          type: 'messages',
          attributes: {
            role: 'user',
            content: content,
            attachments: [],
            createdAt: new Date(),
          },
        };

        patchState(store, (state) => ({
          messages: [...state.messages, userMessage],
          isAgentThinking: true,
        }));
        try {
          const aiMessage: MessageResponse = await lastValueFrom(
            chatService.sendMessage(<string>store.id(), content),
          );

          patchState(store, (state) => ({
            messages: [...state.messages, aiMessage.data],
            isAgentThinking: false,
          }));
        } catch (error: any) {
          let errorText = 'Sorry, I encountered an error while communicating with the server.';

          if (error instanceof HttpErrorResponse && error.status === 422) {
            errorText = error.error?.message || 'Validation error.';

            // Extract the first error message from the 'errors' object if available
            if (error.error?.errors) {
              const firstErrorKey = Object.keys(error.error.errors)[0];
              if (firstErrorKey && error.error.errors[firstErrorKey].length > 0) {
                errorText = error.error.errors[firstErrorKey][0];
              }
            }
          }
          setErrorMessage(errorText, 'assistant');
          patchState(store, () => ({
            isAgentThinking: false,
          }));
        }
      },
      fetchChatHistory: async (chatId: string | null = null) => {
        if (!store.id() && !chatId) {
          setErrorMessage('Please create a new chat session first.', 'assistant');
          return;
        }

        if (chatId) {
          patchState(store, { id: chatId });
        }

        patchState(store, { isLoading: true });

        try {
          const chatHistory = await lastValueFrom(chatService.getMessages(<string>store.id()));
          if (chatHistory.data.length !== 0) {
            patchState(store, {
              messages: chatHistory.data,
              currentPage: chatHistory.meta.current_page,
              lastPage: chatHistory.meta.last_page,
            });
          } else {
            patchState(store, { messages: initialMessageState.messages });
          }
          patchState(store, { isLoading: false });
          if (chatId) {
            await router.navigate(['/', chatId], { replaceUrl: true });
          }
        } catch (error: any) {
          let errorText = 'Failed to load chat history. Please try again.';

          if (error instanceof HttpErrorResponse && error.status === 422) {
            errorText = error.error?.message || 'Validation error.';

            if (error.error?.errors) {
              const firstErrorKey = Object.keys(error.error.errors)[0];
              if (firstErrorKey && error.error.errors[firstErrorKey].length > 0) {
                errorText = error.error.errors[firstErrorKey][0];
              }
            }
          }

          setErrorMessage(errorText, 'assistant');
          patchState(store, { isLoading: false });
        }
      },
      loadOlderMessages: async () => {
        if (
          store.isLoadingPaginated() ||
          store.isLoading() ||
          store.currentPage() >= store.lastPage() ||
          !store.id()
        ) {
          return;
        }

        patchState(store, { isLoadingPaginated: true });
        const nextPage = store.currentPage() + 1;

        try {
          const chatHistory = await lastValueFrom(
            chatService.getMessages(<string>store.id(), nextPage),
          );
          patchState(store, (state) => ({
            messages: [...chatHistory.data, ...state.messages],
            currentPage: chatHistory.meta.current_page,
            lastPage: chatHistory.meta.last_page,
            isLoadingPaginated: false,
          }));
        } catch (err) {
          if (!environment.production) console.error(err);
          patchState(store, { isLoadingPaginated: false });
        }
      },
    };
  }),
);

export const ChatStore = signalStore(
  { providedIn: 'root' },
  withState(initalChatState),
  withComputed((store) => ({
    displayedChats: computed(() => {
      if (!store.searchQuery()) return store._chats();
      return store._filteredChats();
    }),
  })),
  withMethods((store) => {
    const chatService = inject(ChatService);
    return {
      fetchChats: async () => {
        patchState(store, { isLoading: true });
        try {
          const chats = await lastValueFrom(chatService.getChats());
          patchState(store, {
            isLoading: false,
            _chats: chats.data,
          });
        } catch (error: any) {
          patchState(store, { isLoading: false });
        }
      },
      // This method first searches local chats which are loaded on init.
      // If not hit, then it calls backend for results as chats come paginated.
      performSearch: rxMethod<string>(
        pipe(
          debounceTime(300),
          distinctUntilChanged(),
          switchMap(async (query) => {
            patchState(store, { isLoading: true });

            if (!query) {
              patchState(store, { searchQuery: null, _filteredChats: [], isLoading: false });
              return of(null);
            }

            patchState(store, { searchQuery: query });

            const localResults = store._chats().filter((chat) => {
              return chat.attributes.title?.toLowerCase().includes(query.toLowerCase());
            });
            if (localResults.length > 0) {
              patchState(store, {
                _filteredChats: localResults,
              });
            } else {
              const serverResults = await lastValueFrom(chatService.searchChat(query));
              patchState(store, { _serverSearchResults: serverResults.data });
            }
            patchState(store, { isLoading: false });
            return of(null);
          }),
        ),
      ),
    };
  }),
);