import React, { useState, useEffect, useCallback } from 'react' import { Message, ChatSession, MessageStatus } from '@/types/chat' import { chatStorage } from '@/lib/chat-storage' interface UseChatOptions { api_endpoint?: string defaultModel?: string defaultSystemPrompt?: string } interface ApiResponse { thinking_content: string content: string model_used: string supports_thinking: boolean } export function useChat(options: UseChatOptions = {}) { const { api_endpoint = `${window.location.origin}/generate`, defaultModel = 'Qwen/Qwen3-30B-A3B', defaultSystemPrompt = '' } = options // Chat state const [sessions, setSessions] = useState([]) const [currentSessionId, setCurrentSessionId] = useState(null) const [input, setInput] = useState('') const [status, setStatus] = useState({ isLoading: false, error: null }) // Model settings const [selectedModel, setSelectedModel] = useState(defaultModel) const [systemPrompt, setSystemPrompt] = useState(defaultSystemPrompt) const [temperature, setTemperature] = useState(0.7) const [maxTokens, setMaxTokens] = useState(1024) // Current session - add dependency on sessions to force re-render const currentSession = React.useMemo(() => { const session = sessions.find((s: any) => s.id === currentSessionId) || null console.log('useChat - currentSession updated:', { sessionId: currentSessionId, found: !!session, messageCount: session?.messages?.length || 0 }) return session }, [sessions, currentSessionId]) const messages = React.useMemo(() => { const msgs = currentSession?.messages || [] console.log('useChat - messages computed:', msgs.length, msgs.map(m => ({ id: m.id, role: m.role }))) return msgs }, [currentSession?.messages]) // Load sessions on mount useEffect(() => { const loadedSessions = chatStorage.getAllSessions() setSessions(loadedSessions) const currentSession = chatStorage.getCurrentSession() if (currentSession) { setCurrentSessionId(currentSession.id) } else if (loadedSessions.length > 0) { setCurrentSessionId(loadedSessions[0].id) chatStorage.setCurrentSession(loadedSessions[0].id) } }, []) // Create new session const createNewSession = useCallback(() => { const newSession = chatStorage.createSession( undefined, // Auto-generate title selectedModel, systemPrompt ) // Update React state with all sessions from localStorage const updatedSessions = chatStorage.getAllSessions() setSessions([...updatedSessions]) // Force update with new array reference setCurrentSessionId(newSession.id) chatStorage.setCurrentSession(newSession.id) return newSession.id }, [selectedModel, systemPrompt]) // Switch to session const selectSession = useCallback((sessionId: string) => { setCurrentSessionId(sessionId) chatStorage.setCurrentSession(sessionId) }, []) // Delete session const deleteSession = useCallback((sessionId: string) => { chatStorage.deleteSession(sessionId) const updatedSessions = chatStorage.getAllSessions() setSessions([...updatedSessions]) // Force update with new array reference if (currentSessionId === sessionId) { if (updatedSessions.length > 0) { setCurrentSessionId(updatedSessions[0].id) chatStorage.setCurrentSession(updatedSessions[0].id) } else { setCurrentSessionId(null) } } }, [currentSessionId]) // Rename session const renameSession = useCallback((sessionId: string, newTitle: string) => { chatStorage.updateSession(sessionId, { title: newTitle }) const updatedSessions = chatStorage.getAllSessions() setSessions([...updatedSessions]) // Force update with new array reference }, []) // Add message to current session const addMessage = useCallback((message: Omit) => { if (!currentSessionId) return chatStorage.addMessageToSession(currentSessionId, message) // Force update with new array reference const updatedSessions = chatStorage.getAllSessions() setSessions([...updatedSessions]) }, [currentSessionId]) // Send message const sendMessage = useCallback(async () => { if (!input.trim() || status.isLoading) return let sessionId = currentSessionId // Create new session if none exists if (!sessionId) { sessionId = createNewSession() } const userMessage = input.trim() setInput('') setStatus({ isLoading: true, error: null }) // Add user message directly to the specific session if (sessionId) { chatStorage.addMessageToSession(sessionId, { role: 'user', content: userMessage }) // Force update sessions state with fresh data const updatedSessions = chatStorage.getAllSessions() setSessions([...updatedSessions]) // Create new array reference to force re-render } // Add system message if system prompt is set // Check actual session messages from storage, not React state const actualSession = chatStorage.getSession(sessionId!) const hasMessages = actualSession?.messages && actualSession.messages.length > 1 // >1 because user message was just added if (systemPrompt && !hasMessages && sessionId) { chatStorage.addMessageToSession(sessionId, { role: 'system', content: systemPrompt }) // Force update sessions state with fresh data const updatedSessions = chatStorage.getAllSessions() setSessions([...updatedSessions]) // Create new array reference to force re-render } try { // Get conversation history (excluding system messages for the API) const actualSession = chatStorage.getSession(sessionId!) const conversationHistory = actualSession?.messages ?.filter((msg: any) => msg.role !== 'system') ?.map((msg: any) => ({ role: msg.role, content: msg.content })) || [] const response = await fetch(api_endpoint, { method: 'POST', headers: { 'Content-Type': 'application/json', }, body: JSON.stringify({ prompt: userMessage, messages: conversationHistory, system_prompt: systemPrompt || null, model_name: selectedModel, temperature, max_new_tokens: maxTokens }), }) if (!response.ok) { const errorData = await response.json() throw new Error(errorData.detail || `HTTP error! status: ${response.status}`) } const data: ApiResponse = await response.json() // Add assistant message directly to the specific session if (sessionId) { chatStorage.addMessageToSession(sessionId, { role: 'assistant', content: data.content, thinking_content: data.thinking_content, model_used: data.model_used, supports_thinking: data.supports_thinking }) // Force update sessions state with fresh data const updatedSessions = chatStorage.getAllSessions() setSessions([...updatedSessions]) // Create new array reference to force re-render } setStatus({ isLoading: false, error: null }) } catch (error) { const errorMessage = error instanceof Error ? error.message : 'An error occurred' setStatus({ isLoading: false, error: errorMessage }) // Add error message directly to the specific session if (sessionId) { chatStorage.addMessageToSession(sessionId, { role: 'assistant', content: `Sorry, I encountered an error: ${errorMessage}` }) // Force update sessions state with fresh data const updatedSessions = chatStorage.getAllSessions() setSessions([...updatedSessions]) // Create new array reference to force re-render } } }, [ input, status.isLoading, currentSessionId, createNewSession, addMessage, systemPrompt, messages.length, api_endpoint, selectedModel, temperature, maxTokens ]) // Stop generation (placeholder for future implementation) const stopGeneration = useCallback(() => { setStatus({ isLoading: false, error: null }) }, []) // Clear all sessions const clearAllSessions = useCallback(() => { chatStorage.clear() setSessions([]) setCurrentSessionId(null) }, []) return { // Session management sessions, currentSession, currentSessionId, createNewSession, selectSession, deleteSession, renameSession, clearAllSessions, // Messages messages, input, setInput, // Chat actions sendMessage, stopGeneration, // Status isLoading: status.isLoading, error: status.error, // Model settings selectedModel, setSelectedModel, systemPrompt, setSystemPrompt, temperature, setTemperature, maxTokens, setMaxTokens } }