frontend/src/hooks/useChat.ts

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<ChatSession[]>([])
  const [currentSessionId, setCurrentSessionId] = useState<string | null>(null)
  const [input, setInput] = useState('')
  const [status, setStatus] = useState<MessageStatus>({
    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<Message, 'id' | 'timestamp'>) => {
    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
  }
}