nextjs-chat-gpt/components/chat/ChatAgent.tsx

'use client';

import { useState, useCallback, useRef, useEffect } from 'react';
import type { ChatStatus } from 'ai';
import { ChatInput, type PromptInputMessage } from './ChatInput';
import { ChatMessages } from './ChatMessages';
import { ChatPrompt, PromptSuggestion } from './ChatPrompt';
import { ChatSidebar, type Chat } from './ChatSidebar';
import { useChat, UIMessage } from '@ai-sdk/react';
import { DefaultChatTransport } from 'ai';
import {
  createChat,
  deleteChat,
  getMessages,
  listChats,
  saveMessages,
} from '@/services/indexed-db/chat-store';
import {
  SidebarProvider,
  SidebarInset,
  SidebarTrigger,
} from '@/components/ui/sidebar';

interface ChatAgentProps {
  suggestions?: PromptSuggestion[];
  webSearchEnabled?: boolean;
}

async function convertToDataUrl(url: string): Promise<string> {
  const response = await fetch(url);
  if (!response.ok) throw new Error('Could not read the selected file');

  const blob = await response.blob();

  return new Promise((resolve, reject) => {
    const reader = new FileReader();
    reader.onload = () => {
      if (typeof reader.result === 'string') {
        resolve(reader.result);
      } else {
        reject(new Error('Could not encode the selected file'));
      }
    };
    reader.onerror = () => reject(reader.error);
    reader.readAsDataURL(blob);
  });
}

export function ChatAgent({
  suggestions,
  webSearchEnabled = false,
}: ChatAgentProps) {
  const [inputValue, setInputValue] = useState('');
  const [chatId, setChatId] = useState<string | null>(null);
  const [chats, setChats] = useState<Chat[]>([]);
  const chatIdRef = useRef(chatId);
  const pendingChatIdRef = useRef<string | null>(null);
  chatIdRef.current = chatId;

  const fetchChats = useCallback(async () => {
    try {
      setChats(await listChats());
    } catch (error) {
      console.error('Failed to fetch chats:', error);
    }
  }, []);

  const transportRef = useRef(
    new DefaultChatTransport({
      api: '/api/chat',
    })
  );

  const {
    status,
    messages: chatMessages,
    setMessages: setChatMessages,
    sendMessage,
    stop,
    error: aiError,
    addToolApprovalResponse,
  } = useChat({
    transport: transportRef.current,
    onFinish: async ({ messages }) => {
      const activeChatId = pendingChatIdRef.current ?? chatIdRef.current;
      if (!activeChatId) return;

      try {
        await saveMessages(activeChatId, messages);
        await fetchChats();
      } catch (error) {
        console.error('Failed to save messages:', error);
      } finally {
        if (pendingChatIdRef.current === activeChatId) {
          pendingChatIdRef.current = null;
        }
      }
    },
    onError: (error) => {
      console.error('Chat error:', error);
    },
  });

  const [isPreparingFiles, setIsPreparingFiles] = useState(false);

  const chatStatus: ChatStatus = isPreparingFiles
    ? 'submitted'
    : status === 'submitted' || status === 'streaming'
      ? status
      : 'ready';

  useEffect(() => {
    fetchChats();
  }, [fetchChats]);

  const handleSelectChat = useCallback(
    async (selectedChatId: string) => {
      try {
        const messages = await getMessages(selectedChatId);
        chatIdRef.current = selectedChatId;
        setChatId(selectedChatId);

        setChatMessages(messages);
      } catch (error) {
        console.error('Failed to load chat:', error);
      }
    },
    [setChatMessages]
  );

  const handleNewChat = useCallback(() => {
    chatIdRef.current = null;
    setChatId(null);
    setChatMessages([]);
    setInputValue('');
  }, [setChatMessages]);

  const handleDeleteChat = useCallback(
    async (deletedChatId: string) => {
      try {
        await deleteChat(deletedChatId);
        setChats((prev) => prev.filter((c) => c.id !== deletedChatId));
        if (chatId === deletedChatId) {
          chatIdRef.current = null;
          setChatId(null);
          setChatMessages([]);
        }
      } catch (error) {
        console.error('Failed to delete chat:', error);
      }
    },
    [chatId, setChatMessages]
  );

  const handleSendMessage = useCallback(
    async (message: PromptInputMessage) => {
      const trimmedValue = message.text?.trim();
      if (!trimmedValue && !message.files?.length) return;
      if (
        status === 'submitted' ||
        status === 'streaming' ||
        isPreparingFiles
      )
        return;

      let activeChatId = chatIdRef.current;

      if (!chatIdRef.current) {
        const newChatId = crypto.randomUUID();
        const title = trimmedValue ? trimmedValue.slice(0, 100) : 'New Chat';

        activeChatId = newChatId;
        chatIdRef.current = newChatId;
        setChatId(newChatId);

        try {
          await createChat(newChatId, title);
          await fetchChats();
        } catch (error) {
          console.error('Failed to create chat:', error);
        }
      }

      const parts: Array<
        | { type: 'text'; text: string }
        | { type: 'file'; url: string; mediaType: string; filename: string }
      > = [];

      if (trimmedValue) {
        parts.push({ type: 'text', text: trimmedValue });
      }

      // PromptInput supplies files as base64 data URLs. Keep them inline so the
      // demo never needs a remote storage bucket.
      if (message.files?.length) {
        setIsPreparingFiles(true);
        for (const file of message.files) {
          try {
            const dataUrl = file.url.startsWith('data:')
              ? file.url
              : await convertToDataUrl(file.url);

            parts.push({
              type: 'file',
              url: dataUrl,
              mediaType: file.mediaType || 'application/octet-stream',
              filename: file.filename || 'file',
            });
          } catch (err) {
            console.error('File encoding error:', err);
          }
        }
        setIsPreparingFiles(false);
      }

      if (parts.length === 0) return;

      const userMessage: UIMessage = {
        id: crypto.randomUUID(),
        role: 'user',
        parts,
      };

      setInputValue('');

      if (activeChatId) {
        try {
          await saveMessages(activeChatId, [...chatMessages, userMessage]);
        } catch (error) {
          console.error('Failed to save user message:', error);
        }
      }

      pendingChatIdRef.current = activeChatId;
      await sendMessage(userMessage);
    },
    [
      status,
      sendMessage,
      chatMessages,
      isPreparingFiles,
      fetchChats,
    ]
  );

  const handleStop = useCallback(() => {
    stop();
  }, [stop]);

  const hasMessages = chatMessages.length > 0;

  return (
    <SidebarProvider defaultOpen={false}>
      <ChatSidebar
        chats={chats}
        activeChatId={chatId}
        onSelectChat={handleSelectChat}
        onNewChat={handleNewChat}
        onDeleteChat={handleDeleteChat}
      />
      <SidebarInset className="bg-muted">
        {!hasMessages ? (
          <div className="relative">
            <div className="absolute top-2 left-2 z-10">
              <SidebarTrigger />
            </div>
            <ChatPrompt
              value={inputValue}
              onChange={setInputValue}
              onSubmit={handleSendMessage}
              status={chatStatus}
              suggestions={suggestions}
              webSearchEnabled={webSearchEnabled}
            />
          </div>
        ) : (
          <div className="relative flex h-screen flex-col bg-muted">
            <div className="absolute top-2 left-2 z-10">
              <SidebarTrigger />
            </div>
            <div className="flex-1 overflow-hidden">
              <ChatMessages
                messages={chatMessages}
                isLoading={status === 'submitted' || status === 'streaming'}
              />
            </div>
            <div className="sticky bottom-0 bg-muted p-4">
              <div className="mx-auto max-w-3xl">
                <ChatInput
                  value={inputValue}
                  onChange={setInputValue}
                  onSubmit={handleSendMessage}
                  status={chatStatus}
                />
              </div>
            </div>
          </div>
        )}
      </SidebarInset>
    </SidebarProvider>
  );
}

export type { UIMessage as Message };