Momento/memento-note/components/chat/chat-container.tsx

'use client'

import { useState, useEffect, useRef, useCallback } from 'react'
import { useChat } from '@ai-sdk/react'
import { DefaultChatTransport } from 'ai'
import { ChatSidebar } from './chat-sidebar'
import { ChatMessages } from './chat-messages'
import { ChatInput } from './chat-input'
import { createConversation, getConversationDetails, getConversations, deleteConversation } from '@/app/actions/chat-actions'
import { toast } from 'sonner'
import type { UIMessage } from 'ai'
import { useLanguage } from '@/lib/i18n'

interface ChatContainerProps {
  initialConversations: any[]
  notebooks: any[]
  webSearchAvailable?: boolean
}

export function ChatContainer({ initialConversations, notebooks, webSearchAvailable }: ChatContainerProps) {
  const { t, language } = useLanguage()
  const [conversations, setConversations] = useState(initialConversations)
  const [currentId, setCurrentId] = useState<string | null>(null)
  const [selectedNotebook, setSelectedNotebook] = useState<string | undefined>(undefined)
  const [webSearchEnabled, setWebSearchEnabled] = useState(false)
  const [historyMessages, setHistoryMessages] = useState<UIMessage[]>([])
  const [isLoadingHistory, setIsLoadingHistory] = useState(false)

  // Prevents the useEffect from loading an empty conversation
  // when we just created one via createConversation()
  const skipHistoryLoad = useRef(false)
  const scrollRef = useRef<HTMLDivElement>(null)

  const transport = useRef(new DefaultChatTransport({
    api: '/api/chat',
  })).current

  const {
    messages,
    sendMessage,
    status,
    setMessages,
    stop,
  } = useChat({
    transport,
    onError: (error) => {
      toast.error(error.message || t('chat.assistantError'))
    },
  })

  const isLoading = status === 'submitted' || status === 'streaming'

  const refreshConversations = useCallback(async () => {
    try {
      const updated = await getConversations()
      setConversations(updated)
    } catch {}
  }, [])

  // Timeout warning: show toast if response takes > 30s
  useEffect(() => {
    if (!isLoading) return
    const timer = setTimeout(() => {
      toast.warning(t('chat.timeoutWarning') || 'Response is taking longer than expected...')
    }, 30000)
    return () => clearTimeout(timer)
  }, [isLoading, t])

  // Sync historyMessages after each completed streaming response
  // so the display doesn't revert to stale history.
  // Also refresh sidebar so new conversation appears.
  useEffect(() => {
    if (status === 'ready' && messages.length > 0) {
      setHistoryMessages([...messages])
      refreshConversations()
    }
  }, [status, messages, refreshConversations])

  // Load conversation details when the user selects a different conversation
  useEffect(() => {
    // Skip if we just created the conversation — useChat already has the messages
    if (skipHistoryLoad.current) {
      skipHistoryLoad.current = false
      return
    }

    if (currentId) {
      const loadMessages = async () => {
        setIsLoadingHistory(true)
        try {
          const details = await getConversationDetails(currentId)
          if (details) {
            const loaded: UIMessage[] = details.messages.map((m: any, i: number) => ({
              id: m.id || `hist-${i}`,
              role: m.role as 'user' | 'assistant',
              parts: [{ type: 'text' as const, text: m.content }],
            }))
            setHistoryMessages(loaded)
            setMessages(loaded)
          }
        } catch (error) {
          toast.error(t('chat.loadError'))
        } finally {
          setIsLoadingHistory(false)
        }
      }
      loadMessages()
    } else {
      setMessages([])
      setHistoryMessages([])
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [currentId])

  const handleSendMessage = async (content: string, notebookId?: string) => {
    if (notebookId) {
      setSelectedNotebook(notebookId)
    }

    // If no active conversation, create one BEFORE streaming
    let convId = currentId
    if (!convId) {
      try {
        const result = await createConversation(content, notebookId || selectedNotebook)
        convId = result.id
        // Tell the useEffect to skip — we don't want to load an empty conversation
        skipHistoryLoad.current = true
        setCurrentId(convId)
        setHistoryMessages([])
        setConversations((prev) => [
          { id: result.id, title: result.title, updatedAt: new Date() },
          ...prev,
        ])
      } catch {
        toast.error(t('chat.createError'))
        return
      }
    }

    await sendMessage(
      { text: content },
      {
        body: {
          conversationId: convId,
          notebookId: notebookId || selectedNotebook || undefined,
          language,
          webSearch: webSearchEnabled,
        },
      }
    )
  }

  const handleNewChat = () => {
    setCurrentId(null)
    setMessages([])
    setHistoryMessages([])
    setSelectedNotebook(undefined)
    setWebSearchEnabled(false)
  }

  const handleDeleteConversation = async (id: string) => {
    try {
      await deleteConversation(id)
      if (currentId === id) {
        handleNewChat()
      }
      await refreshConversations()
    } catch {
      toast.error(t('chat.deleteError'))
    }
  }

  // During streaming or if useChat has more messages than history, prefer useChat
  const displayMessages = isLoading || messages.length > historyMessages.length
    ? messages
    : historyMessages

  // Auto-scroll to bottom when messages change
  useEffect(() => {
    if (scrollRef.current) {
      scrollRef.current.scrollTop = scrollRef.current.scrollHeight
    }
  }, [displayMessages])

  return (
    <div className="flex-1 flex overflow-hidden bg-white dark:bg-[#1a1c22]">
      <ChatSidebar
        conversations={conversations}
        currentId={currentId}
        onSelect={setCurrentId}
        onNew={handleNewChat}
        onDelete={handleDeleteConversation}
      />

      <div className="flex-1 flex flex-col h-full overflow-hidden">
        <div ref={scrollRef} className="flex-1 overflow-y-auto scrollbar-hide pb-6 w-full flex justify-center">
          <ChatMessages messages={displayMessages} isLoading={isLoading || isLoadingHistory} />
        </div>

        <div className="w-full flex justify-center sticky bottom-0 bg-gradient-to-t from-white dark:from-[#1a1c22] via-white/90 dark:via-[#1a1c22]/90 to-transparent pt-6 pb-4">
          <div className="w-full max-w-4xl px-4">
            <ChatInput
              onSend={handleSendMessage}
              isLoading={isLoading}
              onStop={stop}
              notebooks={notebooks}
              currentNotebookId={selectedNotebook || null}
              webSearchEnabled={webSearchEnabled}
              onToggleWebSearch={() => setWebSearchEnabled(prev => !prev)}
              webSearchAvailable={webSearchAvailable}
            />
          </div>
        </div>
      </div>
    </div>
  )
}