homeai/homeai-dashboard/src/hooks/useChat.js

import { useState, useCallback, useEffect, useRef } from 'react'
import { sendMessage } from '../lib/api'
import { getConversation, saveConversation } from '../lib/conversationApi'

export function useChat(conversationId, conversationMeta, onConversationUpdate) {
  const [messages, setMessages] = useState([])
  const [isLoading, setIsLoading] = useState(false)
  const [isLoadingConv, setIsLoadingConv] = useState(false)
  const convRef = useRef(null)
  const idRef = useRef(conversationId)

  // Keep idRef in sync
  useEffect(() => { idRef.current = conversationId }, [conversationId])

  // Load conversation from server when ID changes
  useEffect(() => {
    if (!conversationId) {
      setMessages([])
      convRef.current = null
      return
    }

    let cancelled = false
    setIsLoadingConv(true)

    getConversation(conversationId).then(conv => {
      if (cancelled) return
      if (conv) {
        convRef.current = conv
        setMessages(conv.messages || [])
      } else {
        convRef.current = null
        setMessages([])
      }
      setIsLoadingConv(false)
    }).catch(() => {
      if (!cancelled) {
        convRef.current = null
        setMessages([])
        setIsLoadingConv(false)
      }
    })

    return () => { cancelled = true }
  }, [conversationId])

  // Persist conversation to server
  const persist = useCallback(async (updatedMessages, title, overrideId) => {
    const id = overrideId || idRef.current
    if (!id) return
    const now = new Date().toISOString()
    const conv = {
      id,
      title: title || convRef.current?.title || '',
      characterId: conversationMeta?.characterId || convRef.current?.characterId || '',
      characterName: conversationMeta?.characterName || convRef.current?.characterName || '',
      createdAt: convRef.current?.createdAt || now,
      updatedAt: now,
      messages: updatedMessages,
    }
    convRef.current = conv
    await saveConversation(conv).catch(() => {})
    if (onConversationUpdate) {
      onConversationUpdate(id, {
        title: conv.title,
        updatedAt: conv.updatedAt,
        messageCount: conv.messages.length,
      })
    }
  }, [conversationMeta, onConversationUpdate])

  // send accepts an optional overrideId for when the conversation was just created
  // and an optional promptStyle to control response style
  const send = useCallback(async (text, overrideId, promptStyle) => {
    if (!text.trim() || isLoading) return null

    const userMsg = { id: Date.now(), role: 'user', content: text.trim(), timestamp: new Date().toISOString() }
    const isFirstMessage = messages.length === 0
    const newMessages = [...messages, userMsg]
    setMessages(newMessages)
    setIsLoading(true)

    try {
      const activeConvId = overrideId || idRef.current
      const { response, model, prompt_style } = await sendMessage(text.trim(), conversationMeta?.characterId || null, promptStyle, activeConvId)
      const assistantMsg = {
        id: Date.now() + 1,
        role: 'assistant',
        content: response,
        timestamp: new Date().toISOString(),
        ...(model && { model }),
        ...(prompt_style && { prompt_style }),
      }
      const allMessages = [...newMessages, assistantMsg]
      setMessages(allMessages)

      const title = isFirstMessage
        ? text.trim().slice(0, 80) + (text.trim().length > 80 ? '...' : '')
        : undefined
      await persist(allMessages, title, overrideId)

      return response
    } catch (err) {
      const errorMsg = {
        id: Date.now() + 1,
        role: 'assistant',
        content: `Error: ${err.message}`,
        timestamp: new Date().toISOString(),
        isError: true,
      }
      const allMessages = [...newMessages, errorMsg]
      setMessages(allMessages)
      await persist(allMessages, undefined, overrideId)
      return null
    } finally {
      setIsLoading(false)
    }
  }, [isLoading, messages, persist])

  // Retry: remove the error message, re-send the preceding user message
  const retry = useCallback(async (errorMsgId, promptStyle) => {
    const idx = messages.findIndex(m => m.id === errorMsgId)
    if (idx < 1) return null
    // Find the user message right before the error
    const userMsg = messages[idx - 1]
    if (!userMsg || userMsg.role !== 'user') return null
    // Remove the error message
    const cleaned = messages.filter(m => m.id !== errorMsgId)
    setMessages(cleaned)
    await persist(cleaned)
    // Re-send (but we need to temporarily set messages back without the error so send picks up correctly)
    // Instead, inline the send logic with the cleaned message list
    setIsLoading(true)
    try {
      const activeConvId = idRef.current
      const { response, model, prompt_style } = await sendMessage(userMsg.content, conversationMeta?.characterId || null, promptStyle, activeConvId)
      const assistantMsg = {
        id: Date.now() + 1,
        role: 'assistant',
        content: response,
        timestamp: new Date().toISOString(),
        ...(model && { model }),
        ...(prompt_style && { prompt_style }),
      }
      const allMessages = [...cleaned, assistantMsg]
      setMessages(allMessages)
      await persist(allMessages)
      return response
    } catch (err) {
      const newError = {
        id: Date.now() + 1,
        role: 'assistant',
        content: `Error: ${err.message}`,
        timestamp: new Date().toISOString(),
        isError: true,
      }
      const allMessages = [...cleaned, newError]
      setMessages(allMessages)
      await persist(allMessages)
      return null
    } finally {
      setIsLoading(false)
    }
  }, [messages, persist, conversationMeta])

  const clearHistory = useCallback(async () => {
    setMessages([])
    if (idRef.current) {
      await persist([], undefined)
    }
  }, [persist])

  return { messages, isLoading, isLoadingConv, send, retry, clearHistory }
}