vercel__workflow/docs/components/geistdocs/chat.tsx

'use client';

import type { UIMessage } from '@ai-sdk/react';
import { useChat } from '@ai-sdk/react';
import { DefaultChatTransport } from 'ai';
import { useLiveQuery } from 'dexie-react-hooks';
import { ChevronRightIcon, Trash } from 'lucide-react';
import { Portal } from 'radix-ui';
import { useCallback, useEffect, useRef, useState } from 'react';
import { toast } from 'sonner';
import type { MyUIMessage } from '@/app/api/chat/types';
import {
  Conversation,
  ConversationContent,
  ConversationScrollButton,
} from '@/components/ai-elements/conversation';
import {
  Message,
  MessageContent,
  MessageResponse,
} from '@/components/ai-elements/message';
import {
  PromptInput,
  PromptInputBody,
  PromptInputFooter,
  type PromptInputProps,
  PromptInputProvider,
  PromptInputSubmit,
  PromptInputTextarea,
} from '@/components/ai-elements/prompt-input';
import { Suggestion, Suggestions } from '@/components/ai-elements/suggestion';
import { Button } from '@/components/ui/button';
import { Drawer, DrawerContent, DrawerTrigger } from '@/components/ui/drawer';
import { useChatContext } from '@/hooks/geistdocs/use-chat';
import { useIsMobile } from '@/hooks/use-mobile';
import { db } from '@/lib/geistdocs/db';
import { cn } from '@/lib/utils';
import { ButtonGroup } from '../ui/button-group';
import { Kbd, KbdGroup } from '../ui/kbd';
import { Spinner } from '../ui/spinner';
import { Tooltip, TooltipContent, TooltipTrigger } from '../ui/tooltip';
import { CopyChat } from './copy-chat';
import { MessageMetadata } from './message-metadata';

const isFromPreviousDay = (timestamp: number): boolean => {
  const messageDate = new Date(timestamp);
  const today = new Date();

  return (
    messageDate.getFullYear() !== today.getFullYear() ||
    messageDate.getMonth() !== today.getMonth() ||
    messageDate.getDate() !== today.getDate()
  );
};

export const useChatPersistence = () => {
  const saveTimeoutRef = useRef<ReturnType<typeof setTimeout> | undefined>(
    undefined
  );

  // Load messages from Dexie with live query
  const storedMessages = useLiveQuery(() =>
    db.messages.orderBy('sequence').toArray()
  );

  // Clear messages if they're from a previous day
  useEffect(() => {
    if (storedMessages && storedMessages.length > 0) {
      const firstMessage = storedMessages[0];
      if (firstMessage && isFromPreviousDay(firstMessage.timestamp)) {
        db.messages.clear();
      }
    }
  }, [storedMessages]);

  // Filter out stale messages from previous days
  const freshMessages = storedMessages?.filter(
    (msg) => !isFromPreviousDay(msg.timestamp)
  );

  const initialMessages =
    freshMessages?.map(({ timestamp, sequence, ...message }) => message) ?? [];

  const isLoading = storedMessages === undefined;

  // Save messages to Dexie with debouncing
  const saveMessages = useCallback((messages: UIMessage[]) => {
    if (saveTimeoutRef.current) {
      clearTimeout(saveTimeoutRef.current);
    }

    saveTimeoutRef.current = setTimeout(async () => {
      try {
        const baseTimestamp = Date.now();
        const messagesToStore = messages.map((message, index) => ({
          ...message,
          timestamp: baseTimestamp + index * 1000,
          sequence: index,
        }));

        await db.transaction('rw', db.messages, async () => {
          await db.messages.clear();
          await db.messages.bulkAdd(messagesToStore);
        });
      } catch (error) {
        console.error('Failed to save messages:', error);
      }
    }, 300);
  }, []);

  // Clear all messages from Dexie
  const clearMessages = useCallback(async () => {
    try {
      await db.messages.clear();
    } catch (error) {
      console.error('Failed to clear messages:', error);
    }
  }, []);

  // Cleanup timeout on unmount
  useEffect(
    () => () => {
      if (saveTimeoutRef.current) {
        clearTimeout(saveTimeoutRef.current);
      }
    },
    []
  );

  return {
    initialMessages,
    isLoading,
    saveMessages,
    clearMessages,
  };
};

interface ChatProps {
  basePath: string | undefined;
  suggestions: string[];
}

type ChatInnerProps = ChatProps & {
  isOpen: boolean;
};

const ChatInner = ({ basePath, suggestions, isOpen }: ChatInnerProps) => {
  const textareaRef = useRef<HTMLTextAreaElement>(null);
  const [isInitialized, setIsInitialized] = useState(false);
  const [localPrompt, setLocalPrompt] = useState('');
  const [providerKey, setProviderKey] = useState(0);
  const { prompt, setPrompt, setIsOpen } = useChatContext();
  const { initialMessages, isLoading, saveMessages, clearMessages } =
    useChatPersistence();

  const { messages, sendMessage, status, setMessages, stop } = useChat({
    transport: new DefaultChatTransport({
      api: basePath ? `${basePath}/api/chat` : '/api/chat',
    }),
    onError: (error) => {
      toast.error(error.message, {
        description: error.message,
      });
    },
  });

  // Sync external prompt changes to local state and force provider remount
  useEffect(() => {
    if (prompt && prompt !== localPrompt) {
      setLocalPrompt(prompt);
      setProviderKey((prev) => prev + 1);
    }
  }, [prompt, localPrompt]);

  // Set initial messages once loaded from IndexedDB
  useEffect(() => {
    if (!(isLoading || isInitialized) && initialMessages.length > 0) {
      setMessages(initialMessages);
      setIsInitialized(true);
    } else if (!(isLoading || isInitialized)) {
      // Mark as initialized even if no messages to avoid infinite re-runs
      setIsInitialized(true);
    }
  }, [isLoading, initialMessages, isInitialized, setMessages]);

  // Save messages to IndexedDB whenever they change (but only after initialization)
  useEffect(() => {
    if (isInitialized && messages.length > 0) {
      saveMessages(messages);
    }
  }, [messages, saveMessages, isInitialized]);

  // Focus textarea when chat opens
  useEffect(() => {
    if (isOpen) {
      // Small delay to ensure the panel/drawer animation has started
      const timer = setTimeout(() => {
        textareaRef.current?.focus();
      }, 100);
      return () => clearTimeout(timer);
    }
  }, [isOpen]);

  const handleSuggestionClick = async (text: string) => {
    if (status === 'streaming' || status === 'submitted') {
      return;
    }
    setLocalPrompt('');
    setPrompt('');
    await sendMessage({ text });
  };

  const handleSubmit: PromptInputProps['onSubmit'] = async (message, event) => {
    event.preventDefault();

    if (status === 'streaming' || status === 'submitted') {
      return;
    }

    const { text } = message;

    if (!text) {
      return;
    }

    setLocalPrompt('');
    setPrompt('');
    await sendMessage({ text });
  };

  const handleClearChat = async () => {
    try {
      await clearMessages();
      setMessages([]);
      toast.success('Chat history cleared');
    } catch (error) {
      toast.error('Failed to clear chat history', {
        description: error instanceof Error ? error.message : 'Unknown error',
      });
    }
  };

  // Show loading state while initial messages are being loaded
  if (isLoading) {
    return (
      <div className="flex size-full w-full flex-col items-center justify-center overflow-hidden whitespace-nowrap rounded-xl xl:max-w-md xl:border xl:bg-background">
        <Spinner />
      </div>
    );
  }

  return (
    <div className="flex size-full w-full flex-col overflow-hidden whitespace-nowrap bg-background">
      <div className="flex items-center justify-between px-4 py-2.5">
        <h2 className="font-semibold text-sm">Chat</h2>
        <div className="flex items-center gap-3">
          <ButtonGroup orientation="horizontal">
            <CopyChat messages={messages} />
            <Tooltip>
              <TooltipTrigger asChild>
                <Button
                  disabled={messages.length === 0}
                  onClick={handleClearChat}
                  size="icon-sm"
                  variant="ghost"
                >
                  <Trash className="size-3.5" />
                </Button>
              </TooltipTrigger>
              <TooltipContent>Clear chat</TooltipContent>
            </Tooltip>
            <Tooltip>
              <TooltipTrigger asChild>
                <Button
                  onClick={() => setIsOpen(false)}
                  size="icon-sm"
                  variant="ghost"
                >
                  <ChevronRightIcon className="size-3.5" />
                </Button>
              </TooltipTrigger>
              <TooltipContent>Close chat</TooltipContent>
            </Tooltip>
          </ButtonGroup>
        </div>
      </div>

      <Conversation>
        <ConversationContent>
          {messages.map((message) => (
            <Message
              className="max-w-[90%]"
              from={message.role}
              key={message.id}
            >
              <MessageMetadata
                inProgress={status === 'submitted' || status === 'streaming'}
                parts={message.parts as MyUIMessage['parts']}
              />
              {message.parts
                .filter((part) => part.type === 'text')
                .map((part) => (
                  <MessageContent
                    key={`${message.id}-${part.type}-${part.text}`}
                  >
                    <MessageResponse className="text-wrap">
                      {part.text}
                    </MessageResponse>
                  </MessageContent>
                ))}
            </Message>
          ))}
          {status === 'submitted' && (
            <div className="size-12 text-muted-foreground text-sm">
              <Spinner />
            </div>
          )}
        </ConversationContent>
        <ConversationScrollButton className="border-none bg-foreground text-background hover:bg-foreground/80 hover:text-background" />
      </Conversation>

      <div className="relative grid w-auto shrink-0 gap-4 p-4">
        {!messages.length && (
          <>
            <Suggestions className="flex-col items-start gap-px">
              {suggestions.map((text) => (
                <Suggestion
                  className="rounded-none p-0"
                  key={text}
                  onClick={handleSuggestionClick}
                  suggestion={text}
                  variant="link"
                />
              ))}
            </Suggestions>
            <p className="text-muted-foreground text-sm">
              Tip: You can open and close chat with{' '}
              <KbdGroup>
                <Kbd className="border bg-transparent">⌘</Kbd>
                <Kbd className="border bg-transparent">I</Kbd>
              </KbdGroup>
            </p>
          </>
        )}
        <PromptInputProvider initialInput={localPrompt} key={providerKey}>
          <PromptInput onSubmit={handleSubmit}>
            <PromptInputBody>
              <PromptInputTextarea
                maxLength={1000}
                onChange={(e) => {
                  setLocalPrompt(e.target.value);
                  setPrompt(e.target.value);
                }}
                ref={textareaRef}
              />
            </PromptInputBody>
            <PromptInputFooter>
              <p className="text-muted-foreground text-xs">
                {localPrompt.length} / 1000
              </p>
              <PromptInputSubmit
                onClick={
                  status === 'streaming'
                    ? (e) => {
                        e.preventDefault();
                        stop();
                      }
                    : undefined
                }
                status={status}
              />
            </PromptInputFooter>
          </PromptInput>
        </PromptInputProvider>
      </div>
    </div>
  );
};

export const Chat = ({ basePath, suggestions }: ChatProps) => {
  const { isOpen, setIsOpen } = useChatContext();
  const isMobile = useIsMobile();

  useEffect(() => {
    const handleKeyDown = (event: KeyboardEvent) => {
      // Check for Meta (⌘ on Mac, Windows key on Windows) + "i" (ignore case)
      if (
        (event.metaKey || event.ctrlKey) &&
        !event.altKey &&
        !event.shiftKey &&
        event.key.toLowerCase() === 'i'
      ) {
        event.preventDefault();

        setIsOpen((prev) => !prev);
      }
    };

    window.addEventListener('keydown', handleKeyDown);

    return () => {
      window.removeEventListener('keydown', handleKeyDown);
    };
  }, [setIsOpen]);

  return (
    <>
      <Button
        className="hidden shrink-0 shadow-none md:flex"
        onClick={() => setIsOpen(!isOpen)}
        size="sm"
        variant="outline"
      >
        Ask AI
      </Button>

      <Portal.Root className="hidden md:block">
        <div
          className={cn(
            'fixed z-50 flex flex-col gap-4 bg-background transition-all',
            'inset-y-0 right-0 h-full w-3/4 border-l sm:max-w-sm',
            'translate-x-full data-[state=open]:translate-x-0'
          )}
          data-state={isOpen ? 'open' : 'closed'}
        >
          <ChatInner
            basePath={basePath}
            isOpen={isOpen}
            suggestions={suggestions}
          />
        </div>
      </Portal.Root>
      <div className="md:hidden">
        <Drawer
          onOpenChange={isMobile ? setIsOpen : undefined}
          open={isMobile ? isOpen : false}
        >
          <DrawerTrigger asChild>
            <Button className="shadow-none" size="sm" variant="outline">
              Ask AI
            </Button>
          </DrawerTrigger>
          <DrawerContent className="h-[80dvh]">
            <ChatInner
              basePath={basePath}
              isOpen={isOpen}
              suggestions={suggestions}
            />
          </DrawerContent>
        </Drawer>
      </div>
    </>
  );
};