import { useAtom, useSetAtom } from 'jotai';
import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
import { messagesAtom, addMessageAtom, updateChatTitleAtom } from '@/atoms/chatAtoms';

import { sendMessageChatId, subscribeToChat } from '@/hooks/useChatWebsocket';
import type { Payload } from '@/hooks/useChatWebsocket';
import { useComposer } from './useComposer';

export function useChat(chatId: string) {
  const [isLoading, setIsLoading] = useState<boolean>(false);

  const messagesAtomForChat = useMemo(() => messagesAtom(chatId), [chatId]);
  const [messages, setMessages] = useAtom(messagesAtomForChat);
  const { setIdle } = useComposer();

  const addMessage = useSetAtom(addMessageAtom);
  const updateChatTitle = useSetAtom(updateChatTitleAtom);

  const streamingMessageIdRef = useRef<string | null>(null);

  useEffect(() => {
    const unsubscribe = subscribeToChat(chatId, (payload: Payload) => {
      if (payload.action === 'done') {
        setIdle();
        streamingMessageIdRef.current = null;
        return;
      }

      if (payload.action === 'title_updated' && streamingMessageIdRef.current) {
        updateChatTitle({ chatId, title: payload.title });
      }

      if (payload.action === 'image' && payload.url && streamingMessageIdRef.current) {
        setMessages(prev =>
          prev.map(m =>
            m.id === streamingMessageIdRef.current
              ? { ...m, image_url: payload.url }
              : m,
          ),
        );
      }
      
      if (payload.action === 'append' && payload.content && streamingMessageIdRef.current) {
        setMessages(prev =>
          prev.map(m =>
            m.id === streamingMessageIdRef.current
              ? { ...m, content: m.content + payload.content }
              : m,
          ),
        );
      }
    });
    
    return unsubscribe;
  }, [chatId, setIdle, setMessages]);

  const sendUserMessage = useCallback(
    (content: string, overrideChatID?: string) => {
      if (!(chatId || overrideChatID) || !content.trim()) return;
      
      const effectiveChatId = (chatId === 'new' || !chatId) && overrideChatID ? overrideChatID : chatId;
      
      addMessage({
        chatId: effectiveChatId,
        message: { role: 'user', content: content.trim(), type: 'text' },
      });
      
      const id = crypto.randomUUID();
      streamingMessageIdRef.current = id;
      
      addMessage({
        chatId: effectiveChatId,
        message: { id, role: 'assistant', content: '', type: 'text' },
      });
      
      sendMessageChatId(content, effectiveChatId);
    },
    [chatId, addMessage],
  );

  const stopStreaming = useCallback(() => {
    setIsLoading(false);
    streamingMessageIdRef.current = null;
  }, []);

  return {
    messages: messages ?? [],
    sendUserMessage,
    setMessages,
    isLoading,
    stopStreaming,
  };
}