import { atom } from 'jotai';
import type { SetStateAction } from 'jotai';
import { atomFamily } from 'jotai/utils';

export type MessageRole = 'user' | 'assistant' | 'system';

export type Message = {
  id: string;
  role: MessageRole;
  content: string;
  image_url?: string;
  type: 'text';
};

export type Chat = {
  id: string;
  title: string;
  messages: Message[];
  created_at: number;
};

// Global 
export const chatsAtom = atom<Record<string, Chat>>({});
export const currentChatId = atom<string | null>(null);

export const messagesAtom = atomFamily((chatId: string) =>
  atom(
    (get) => get(chatsAtom)[chatId]?.messages ?? [],
    (get, set, update: SetStateAction<Message[]>) => {
      const chats = get(chatsAtom);
      const prevMessages = chats[chatId]?.messages ?? [];
      const newMessages =
        typeof update === 'function'
          ? (update as (prev: Message[]) => Message[])(prevMessages)
          : update;
      
      const updatedChat: Chat = {
        id: chatId,
        title: chats[chatId]?.title ?? 'New Chat',
        messages: newMessages,
        created_at: chats[chatId]?.created_at ?? new Date(),
      };

      let newChats = { ...chats };
      if (!chats[chatId]) {
        newChats = {
          [chatId]: updatedChat,
          ...chats,
        };
      } else {
        newChats[chatId] = updatedChat;
      }
      set(chatsAtom, newChats);
    },
  ),
);

export const addMessageAtom = atom(
  null,
  (get, set, { chatId, message }: { chatId: string; message: Omit<Message, 'id' | 'created_at'> | Message }) => {
    const newMessage: Message = {
      id: crypto.randomUUID(),
      ...message,
    };
    
    const msgAtom = messagesAtom(chatId);
    const currentMessages = get(msgAtom);
    set(msgAtom, [...currentMessages, newMessage]);
  },
);

export const updateChatTitleAtom = atom(
  null,
  (get, set, { chatId, title }: { chatId: string; title: string }) => {
    const chats = get(chatsAtom);
    const chat = chats[chatId];

    if (!chat) return;
    
    set(chatsAtom, {
      ...chats,
      [chatId]: {
        ...chat,
        title,
      },
    });
  },
);