nextjs-chat-gpt/hooks/use-chat-session.ts

'use client';

import { useState, useCallback, useRef, useEffect } from 'react';
import { useChat } from '@ai-sdk/react';
import { DefaultChatTransport, type ChatStatus, type UIMessage } from 'ai';
import { DEFAULT_OPENROUTER_MODEL_ID, type OpenRouterModelId } from '@/lib/openrouter';
import type { PromptInputMessage } from '@/components/ai-elements/prompt-input';
import {
  createChat,
  deleteChat,
  getMessages,
  listChats,
  saveMessages,
  type StoredChat,
} from '@/services/indexed-db/chat-store';

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 useChatSession(defaultModelId: OpenRouterModelId = DEFAULT_OPENROUTER_MODEL_ID) {
  const [modelId, setModelId] = useState<OpenRouterModelId>(defaultModelId);
  const [inputValue, setInputValue] = useState('');
  const [chatId, setChatId] = useState<string | null>(null);
  const [chats, setChats] = useState<StoredChat[]>([]);
  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,
  } = 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, { body: { modelId } });
    },
    [
      status,
      sendMessage,
      chatMessages,
      isPreparingFiles,
      fetchChats,
      modelId,
    ]
  );

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

  return {
    inputValue,
    setInputValue,
    modelId,
    setModelId,
    chatId,
    chats,
    chatMessages,
    status,
    chatStatus,
    aiError,
    isPreparingFiles,
    handleSelectChat,
    handleNewChat,
    handleDeleteChat,
    handleSendMessage,
    handleStop,
  };
}