import {
  ChatInput,
  ChatMessage,
  ChatMessages,
  ChatSection,
  useChatUI,
  useFile,
} from '@llamaindex/chat-ui'
import { useChat } from 'ai/react'

export function CustomChat() {
  const handler = useChat()
  const { imageUrl, getAnnotations, uploadFile, reset } = useFile({
    uploadAPI: '/chat/upload',
  })
  const annotations = getAnnotations()
  const handleUpload = async (file: File) => {
    try {
      await uploadFile(file)
    } catch (error) {
      console.error(error)
    }
  }
  return (
    <ChatSection
      handler={handler}
      className="mx-auto h-screen max-w-3xl overflow-hidden"
    >
      <CustomChatMessages />
      <ChatInput annotations={annotations} resetUploadedFiles={reset}>
        <div>
          {imageUrl ? (
            <img
              className="max-h-[100px] object-contain"
              src={imageUrl}
              alt="uploaded"
            />
          ) : null}
        </div>
        <ChatInput.Form>
          <ChatInput.Field />
          <ChatInput.Upload
            allowedExtensions={['jpg', 'png', 'jpeg']}
            onUpload={handleUpload}
          />
          <ChatInput.Submit />
        </ChatInput.Form>
      </ChatInput>
    </ChatSection>
  )
}

function CustomChatMessages() {
  const { messages, isLoading, append } = useChatUI()
  return (
    <ChatMessages>
      <ChatMessages.List>
        {messages.map((message, index) => (
          <ChatMessage
            key={index}
            message={message}
            isLast={index === messages.length - 1}
            className="items-start"
          >
            <ChatMessage.Avatar>
              <img
                className="border-1 rounded-full border-[#e711dd]"
                alt="LlamaIndex"
                src="/llama.png"
              />
            </ChatMessage.Avatar>
            <ChatMessage.Content isLoading={isLoading} append={append}>
              <ChatMessage.Content.Image />
              <ChatMessage.Content.Markdown />
              <ChatMessage.Content.DocumentFile />
            </ChatMessage.Content>
            <ChatMessage.Actions />
          </ChatMessage>
        ))}
      </ChatMessages.List>
    </ChatMessages>
  )
}