Chat container component

ChatContainer.tsx

import {
  Box,
  FileUpload,
  Flex,
  Input,
  InputGroup,
  Text,
  VStack,
} from '@chakra-ui/react'
import { useCallback, useLayoutEffect, useMemo, useRef, useState } from 'react'
import { useParams } from '@tanstack/react-router'
import { LuSearch } from 'react-icons/lu'
import useInfiniteScroll from 'react-infinite-scroll-hook'
import { useAuthQuery } from '@/api/auth-query'
import { useProjects } from '@/hooks/project'
import { useDebounce } from '@/hooks/use-debounce'
import { useFetchChatMessagesInfinite } from '@/hooks/chat'
import { SkeletonText } from '@/components/shared/SkeletonText'
import { ChatInput } from '../ChatInput'
import { ChatMessage } from '../ChatMessage'

export const ChatContainer = () => {
  const { data: profile } = useAuthQuery()
  const { data: projects } = useProjects()
  const { projectId } = useParams({ strict: false })
  const currentProject = projects?.find(
    (project) => String(project.id) === String(projectId),
  )

  const [search, setSearch] = useState('')
  const debouncedSearchValue = useDebounce(search.trim(), 300)

  const {
    data: fetchedMessagesInfinite,
    isFetchingPreviousPage,
    hasPreviousPage,
    fetchPreviousPage,
    error,
    isLoading,
  } = useFetchChatMessagesInfinite({
    chatRoomId: currentProject?.chat_room_id ?? '',
    direction: 'before',
    limit: 20,
    search: debouncedSearchValue || void 0,
  })

  const messagesInfinite = useMemo(
    () =>
      [
        ...(fetchedMessagesInfinite?.pages.flatMap((page) => page.results) ??
          []),
      ].sort(
        (a, b) =>
          new Date(a.created_at).getTime() - new Date(b.created_at).getTime(),
      ),
    [fetchedMessagesInfinite],
  )

  const [infiniteRef, { rootRef }] = useInfiniteScroll({
    loading: isFetchingPreviousPage,
    hasNextPage: hasPreviousPage,
    onLoadMore: fetchPreviousPage,
    disabled: Boolean(error),
  })

  const scrollableRootRef = useRef<HTMLDivElement | null>(null)
  const lastScrollDistanceToBottomRef = useRef(0)
  const previousSearchRef = useRef('')

  useLayoutEffect(() => {
    const scrollableRoot = scrollableRootRef.current
    const lastScrollDistanceToBottom = lastScrollDistanceToBottomRef.current
    if (scrollableRoot) {
      scrollableRoot.scrollTop =
        scrollableRoot.scrollHeight - lastScrollDistanceToBottom
    }
  }, [messagesInfinite, rootRef])

  useLayoutEffect(() => {
    const scrollableRoot = scrollableRootRef.current
    const hadSearchBefore = previousSearchRef.current.length > 0
    const hasSearchNow = debouncedSearchValue.length > 0

    if (scrollableRoot && hadSearchBefore && !hasSearchNow) {
      scrollableRoot.scrollTop = scrollableRoot.scrollHeight
    }

    previousSearchRef.current = debouncedSearchValue
  }, [debouncedSearchValue])

  const rootRefSetter = useCallback(
    (node: HTMLDivElement) => {
      rootRef(node)
      scrollableRootRef.current = node
    },
    [rootRef],
  )

  const handleRootScroll = useCallback(() => {
    const rootNode = scrollableRootRef.current
    if (rootNode) {
      const scrollDistanceToBottom = rootNode.scrollHeight - rootNode.scrollTop
      lastScrollDistanceToBottomRef.current = scrollDistanceToBottom
    }
  }, [])

  return (
    <Box w="full" h="full" py={2}>
      {/* search */}
      <InputGroup
        borderRadius="lg"
        startElement={<LuSearch color="var(--chakra-colors-muted-sage)" />}
        mb={1}
        boxShadow="0 2px 8px rgba(0, 0, 0, 0.06)"
      >
        <Input
          borderRadius="lg"
          placeholder="Search..."
          bg="var(--chakra-colors-white)"
          value={search}
          onChange={(e) => {
            lastScrollDistanceToBottomRef.current = 0
            setSearch(e.target.value)
          }}
          css={{
            '&:focus': {
              outline: 'none',
              boxShadow: '0 2px 12px rgba(0, 0, 0, 0.1)',
            },
          }}
        />
      </InputGroup>

      {/* messages */}
      <VStack w="full" h="full" gap={2}>
        <Flex
          flex={1}
          direction="column"
          minH={0}
          overflowY="auto"
          w="full"
          h="full"
          ref={rootRefSetter}
          onScroll={handleRootScroll}
        >
          {hasPreviousPage && (
            <Flex
              ref={infiniteRef}
              my={2}
              mx={4}
              justifyContent="center"
              align="center"
            >
              <SkeletonText text="Loading..." />
            </Flex>
          )}
          <>
            {isLoading ? (
              <Flex align="center" justify="center" flex={1}>
                <SkeletonText text="Loading..." />
              </Flex>
            ) : messagesInfinite.length > 0 && profile ? (
              <VStack gap={5} flex={1} pr="1.5" w="fulll">
                {messagesInfinite.map((message, index) => (
                  <ChatMessage
                    key={`${message.id}-${index}`}
                    message={message}
                    profile={profile}
                    search={search}
                  />
                ))}
              </VStack>
            ) : (
              <VStack gap={2} flex={1} align="center" justify="center">
                <Text
                  color="var(--chakra-colors-muted-sage)"
                  fontSize="sm"
                  fontStyle="italic"
                >
                  {debouncedSearchValue
                    ? 'No messages found'
                    : 'No messages yet'}
                </Text>
                <Text color="var(--chakra-colors-muted-sage)" fontSize="xs">
                  {debouncedSearchValue
                    ? 'Try a different search term'
                    : 'Start the conversation!'}
                </Text>
              </VStack>
            )}
          </>
        </Flex>

        {/* input */}
        <FileUpload.Root
          bg="var(--chakra-colors-white)"
          borderTopLeftRadius="lg"
          borderTopRightRadius="lg"
          flexShrink={0}
        >
          <ChatInput placeholder="Use @ai to ask the assistant" />
        </FileUpload.Root>
      </VStack>
    </Box>
  )
}