Messenger/web/src/store/chatStore.ts

import { create } from "zustand";
import { getChats, getMessages, updateMessageStatus } from "../api/chats";
import type { Chat, DeliveryStatus, Message, MessageType } from "../chat/types";

interface ChatState {
  chats: Chat[];
  activeChatId: number | null;
  messagesByChat: Record<number, Message[]>;
  typingByChat: Record<number, number[]>;
  replyToByChat: Record<number, Message | null>;
  unreadBoundaryByChat: Record<number, number>;
  loadChats: (query?: string) => Promise<void>;
  setActiveChatId: (chatId: number | null) => void;
  loadMessages: (chatId: number) => Promise<void>;
  prependMessage: (chatId: number, message: Message) => void;
  addOptimisticMessage: (params: {
    chatId: number;
    senderId: number;
    type: MessageType;
    text: string | null;
    clientMessageId: string;
  }) => void;
  confirmMessageByClientId: (chatId: number, clientMessageId: string, message: Message) => void;
  removeOptimisticMessage: (chatId: number, clientMessageId: string) => void;
  setMessageDeliveryStatus: (chatId: number, messageId: number, status: DeliveryStatus) => void;
  removeMessage: (chatId: number, messageId: number) => void;
  restoreMessages: (chatId: number, messages: Message[]) => void;
  clearChatMessages: (chatId: number) => void;
  incrementUnread: (chatId: number) => void;
  clearUnread: (chatId: number) => void;
  setTypingUsers: (chatId: number, userIds: number[]) => void;
  setReplyToMessage: (chatId: number, message: Message | null) => void;
  updateChatPinnedMessage: (chatId: number, pinnedMessageId: number | null) => void;
}

export const useChatStore = create<ChatState>((set, get) => ({
  chats: [],
  activeChatId: null,
  messagesByChat: {},
  typingByChat: {},
  replyToByChat: {},
  unreadBoundaryByChat: {},
  loadChats: async (query) => {
    const chats = await getChats(query);
    const currentActive = get().activeChatId;
    const nextActive = chats.some((chat) => chat.id === currentActive) ? currentActive : (chats[0]?.id ?? null);
    set({ chats, activeChatId: nextActive });
  },
  setActiveChatId: (chatId) => set({ activeChatId: chatId }),
  loadMessages: async (chatId) => {
    const unreadCount = get().chats.find((c) => c.id === chatId)?.unread_count ?? 0;
    const messages = await getMessages(chatId);
    const sorted = [...messages].reverse();
    set((state) => ({
      messagesByChat: {
        ...state.messagesByChat,
        [chatId]: sorted
      },
      unreadBoundaryByChat: {
        ...state.unreadBoundaryByChat,
        [chatId]: unreadCount
      },
      chats: state.chats.map((chat) => (chat.id === chatId ? { ...chat, unread_count: 0 } : chat))
    }));
    const lastMessage = sorted[sorted.length - 1];
    if (lastMessage) {
      void updateMessageStatus(chatId, lastMessage.id, "message_read");
    }
  },
  prependMessage: (chatId, message) => {
    const old = get().messagesByChat[chatId] ?? [];
    if (old.some((m) => m.id === message.id || (message.client_message_id && m.client_message_id === message.client_message_id))) {
      return;
    }
    set((state) => ({
      messagesByChat: { ...state.messagesByChat, [chatId]: [...old, message] }
    }));
  },
  addOptimisticMessage: ({ chatId, senderId, type, text, clientMessageId }) => {
    const old = get().messagesByChat[chatId] ?? [];
    if (old.some((m) => m.client_message_id === clientMessageId)) {
      return;
    }
    const now = new Date().toISOString();
    const optimistic: Message = {
      id: -(Date.now() + Math.floor(Math.random() * 10000)),
      chat_id: chatId,
      sender_id: senderId,
      type,
      text,
      created_at: now,
      updated_at: now,
      client_message_id: clientMessageId,
      delivery_status: "sending",
      is_pending: true
    };
    set((state) => ({
      messagesByChat: { ...state.messagesByChat, [chatId]: [...old, optimistic] }
    }));
  },
  confirmMessageByClientId: (chatId, clientMessageId, message) => {
    const old = get().messagesByChat[chatId] ?? [];
    const idx = old.findIndex((m) => m.client_message_id === clientMessageId);
    if (idx === -1) {
      if (!old.some((m) => m.id === message.id)) {
        set((state) => ({
          messagesByChat: {
            ...state.messagesByChat,
            [chatId]: [...old, { ...message, client_message_id: clientMessageId, delivery_status: "sent", is_pending: false }]
          }
        }));
      }
      return;
    }
    const next = [...old];
    next[idx] = {
      ...message,
      client_message_id: clientMessageId,
      delivery_status: "sent",
      is_pending: false
    };
    set((state) => ({
      messagesByChat: { ...state.messagesByChat, [chatId]: next }
    }));
  },
  removeOptimisticMessage: (chatId, clientMessageId) => {
    const old = get().messagesByChat[chatId] ?? [];
    set((state) => ({
      messagesByChat: {
        ...state.messagesByChat,
        [chatId]: old.filter((m) => m.client_message_id !== clientMessageId)
      }
    }));
  },
  setMessageDeliveryStatus: (chatId, messageId, status) => {
    const old = get().messagesByChat[chatId] ?? [];
    const idx = old.findIndex((m) => m.id === messageId);
    if (idx === -1) {
      return;
    }
    const current = old[idx];
    const order: Record<DeliveryStatus, number> = { sending: 1, sent: 2, delivered: 3, read: 4 };
    const currentStatus = current.delivery_status ?? "sent";
    if (order[status] <= order[currentStatus]) {
      return;
    }
    const next = [...old];
    next[idx] = { ...current, delivery_status: status, is_pending: false };
    set((state) => ({
      messagesByChat: { ...state.messagesByChat, [chatId]: next }
    }));
  },
  removeMessage: (chatId, messageId) => {
    const old = get().messagesByChat[chatId] ?? [];
    set((state) => ({
      messagesByChat: {
        ...state.messagesByChat,
        [chatId]: old.filter((m) => m.id !== messageId)
      }
    }));
  },
  restoreMessages: (chatId, messages) => {
    if (!messages.length) {
      return;
    }
    const old = get().messagesByChat[chatId] ?? [];
    const byId = new Map<number, Message>();
    for (const message of old) {
      byId.set(message.id, message);
    }
    for (const message of messages) {
      byId.set(message.id, message);
    }
    const merged = [...byId.values()].sort((a, b) => a.id - b.id);
    set((state) => ({
      messagesByChat: {
        ...state.messagesByChat,
        [chatId]: merged
      }
    }));
  },
  clearChatMessages: (chatId) =>
    set((state) => ({
      messagesByChat: {
        ...state.messagesByChat,
        [chatId]: []
      },
      unreadBoundaryByChat: {
        ...state.unreadBoundaryByChat,
        [chatId]: 0
      },
      chats: state.chats.map((chat) => (chat.id === chatId ? { ...chat, unread_count: 0 } : chat))
    })),
  incrementUnread: (chatId) =>
    set((state) => ({
      chats: state.chats.map((chat) =>
        chat.id === chatId ? { ...chat, unread_count: (chat.unread_count ?? 0) + 1 } : chat
      )
    })),
  clearUnread: (chatId) =>
    set((state) => ({
      chats: state.chats.map((chat) => (chat.id === chatId ? { ...chat, unread_count: 0 } : chat)),
      unreadBoundaryByChat: {
        ...state.unreadBoundaryByChat,
        [chatId]: 0
      }
    })),
  setTypingUsers: (chatId, userIds) =>
    set((state) => ({ typingByChat: { ...state.typingByChat, [chatId]: userIds } })),
  setReplyToMessage: (chatId, message) =>
    set((state) => ({ replyToByChat: { ...state.replyToByChat, [chatId]: message } })),
  updateChatPinnedMessage: (chatId, pinnedMessageId) =>
    set((state) => ({
      chats: state.chats.map((chat) => (chat.id === chatId ? { ...chat, pinned_message_id: pinnedMessageId } : chat))
    }))
}));