import type { ChatRoom, Message, MessageReaction, User } from '@/types';
import { create } from 'zustand';
import { subscribeWithSelector } from 'zustand/middleware';

type ChatState = {
  chatRooms: ChatRoom[];
  selectedRoomId: string | null;
  messages: Record<string, Message[]>;
  onlineUsers: Set<string>;
  typingUsers: Record<string, Set<string>>;

  setChatRooms: (rooms: ChatRoom[]) => void;
  addChatRoom: (room: ChatRoom) => void;
  updateChatRoom: (room: ChatRoom) => void;
  removeChatRoom: (roomId: string) => void;
  setSelectedRoom: (roomId: string | null) => void;
  setMessages: (roomId: string, messages: Message[]) => void;
  addMessage: (message: Message) => void;
  removeMessage: (messageId: string, roomId: string) => void;
  updateMessageReactions: (
    messageId: string,
    reactions: MessageReaction[]
  ) => void;
  setOnlineUsers: (users: Set<string>) => void;
  addOnlineUser: (userId: string) => void;
  removeOnlineUser: (userId: string) => void;
  addTypingUser: (roomId: string, username: string) => void;
  removeTypingUser: (roomId: string, username: string) => void;
  updateMessageSeen: (messageId: string, seenBy: string[]) => void;
  updateUserInRooms: (userId: string, userData: Partial<User>) => void;
};

export const useChatStore = create<ChatState>()(
  subscribeWithSelector((set) => ({
    chatRooms: [],
    selectedRoomId: null,
    messages: {},
    onlineUsers: new Set(),
    typingUsers: {},

    setChatRooms: (rooms) => {
      set({ chatRooms: rooms });
    },

    addChatRoom: (room) => {
      set((state) => ({
        chatRooms: [room, ...state.chatRooms],
        selectedRoomId: room.id,
      }));
    },

    setSelectedRoom: (roomId) => {
      set({ selectedRoomId: roomId });
    },

    setMessages: (roomId, messages) => {
      set((state) => ({
        messages: {
          ...state.messages,
          [roomId]: messages,
        },
      }));
    },

    addMessage: (message) => {
      set((state) => ({
        messages: {
          ...state.messages,
          [message.roomId]: [
            ...(state.messages[message.roomId] || []),
            message,
          ],
        },
      }));
    },

    removeMessage: (messageId, roomId) => {
      set((state) => {
        if (!state.messages[roomId]) return state;

        return {
          messages: {
            ...state.messages,
            [roomId]: state.messages[roomId].filter(
              (message) => message.id !== messageId
            ),
          },
        };
      });
    },

    updateMessageReactions: (messageId, reactions) => {
      set((state) => {
        const newMessages = { ...state.messages };

        Object.keys(newMessages).forEach((roomId) => {
          newMessages[roomId] = newMessages[roomId].map((message) =>
            message.id === messageId ? { ...message, reactions } : message
          );
        });

        return { messages: newMessages };
      });
    },

    setOnlineUsers: (users) => {
      set({ onlineUsers: users });
    },

    addOnlineUser: (userId) => {
      set((state) => ({
        onlineUsers: new Set([...state.onlineUsers, userId]),
      }));
    },

    removeOnlineUser: (userId) => {
      set((state) => {
        const newOnlineUsers = new Set(state.onlineUsers);
        newOnlineUsers.delete(userId);
        return { onlineUsers: newOnlineUsers };
      });
    },

    addTypingUser: (roomId, username) => {
      set((state) => {
        const roomTypingUsers = state.typingUsers[roomId] ?? new Set();
        roomTypingUsers.add(username);

        return {
          typingUsers: {
            ...state.typingUsers,
            [roomId]: roomTypingUsers,
          },
        };
      });
    },

    removeTypingUser: (roomId, username) => {
      set((state) => {
        if (!state.typingUsers[roomId])
          return {
            typingUsers: state.typingUsers,
          };

        const roomTypingUsers = new Set(state.typingUsers[roomId]);
        roomTypingUsers.delete(username);

        return {
          typingUsers: {
            ...state.typingUsers,
            [roomId]: roomTypingUsers,
          },
        };
      });
    },

    updateMessageSeen: (messageId: string, seenBy: string[]) => {
      set((state) => {
        const newMessages = { ...state.messages };

        Object.keys(newMessages).forEach((roomId) => {
          newMessages[roomId] = newMessages[roomId].map((message) =>
            message.id === messageId ? { ...message, seenBy } : message
          );
        });

        return { messages: newMessages };
      });
    },

    updateUserInRooms: (userId: string, userData: Partial<User>) => {
      set((state) => {
        const updatedRooms = state.chatRooms.map((room) => ({
          ...room,
          members: room.members.map((member) =>
            member.userId === userId
              ? {
                  ...member,
                  user: { ...member.user, ...userData },
                }
              : member
          ),
        }));

        return { chatRooms: updatedRooms };
      });
    },

    updateChatRoom: (room) => {
      set((state) => ({
        chatRooms: state.chatRooms.map((r) => (r.id === room.id ? room : r)),
      }));
    },

    removeChatRoom: (roomId) => {
      set((state) => {
        const newMessages = { ...state.messages };
        delete newMessages[roomId];

        return {
          chatRooms: state.chatRooms.filter((r) => r.id !== roomId),
          selectedRoomId:
            state.selectedRoomId === roomId ? null : state.selectedRoomId,
          messages: newMessages,
          onlineUsers: state.onlineUsers,
          typingUsers: state.typingUsers,
        };
      });
    },
  }))
);
