import { InfiniteData } from '@tanstack/react-query'
import { chatQueryClient } from '../../contexts/api_provider'
import { ApiCollection, MessageResource } from '../../types'
import { getMessagesQueryKey } from '../../utils/request/get_messages'
import { updateOrCreateRecordInPagesData } from './page_mutations'

export function optimisticallyUpdateMessage({
  conversationId,
  message,
  text,
  replyRootId,
}: {
  conversationId: number
  message: MessageResource
  text: string
  replyRootId?: string | null
}) {
  type QueryData = InfiniteData<ApiCollection<MessageResource>>

  const queryKey = getMessagesQueryKey({
    conversation_id: conversationId,
    reply_root_id: replyRootId,
  })

  let optimisticMessage = message

  chatQueryClient.setQueryData<QueryData>(queryKey, data =>
    updateOrCreateRecordInPagesData({
      data,
      record: message,
      processRecord: (next, prev) => {
        optimisticMessage = {
          ...(prev || next),
          text,
          pending: true,
        }
        return optimisticMessage
      },
    })
  )

  return optimisticMessage
}