PicoBot/web/src/hooks/useChat.ts

import { useState, useCallback, useRef } from 'react'
import type {
  Command,
  ChatMessage,
  Topic,
  WsOutbound,
  AssistantResponse,
  ToolCall,
  ToolResult,
  ToolPending,
  SessionEstablished,
  SessionCreated,
  SessionList,
} from '../types/protocol'

interface UseChatReturn {
  messages: ChatMessage[]
  currentSessionId: string | null
  currentTopicId: string | null
  topics: Topic[]
  isLoading: boolean
  handleMessage: (content: string) => void
  handleCommand: (command: Command) => void
  clearMessages: () => void
  handleServerMessage: (message: WsOutbound) => void
}

export function useChat(): UseChatReturn {
  const [messages, setMessages] = useState<ChatMessage[]>([])
  const [currentSessionId, setCurrentSessionId] = useState<string | null>(null)
  const [currentTopicId, setCurrentTopicId] = useState<string | null>(null)
  const [topics, setTopics] = useState<Topic[]>([])
  const [isLoading, setIsLoading] = useState(false)

  // Message ID generator
  const messageIdCounter = useRef(0)
  const generateMessageId = () => {
    messageIdCounter.current += 1
    return `msg_${Date.now()}_${messageIdCounter.current}`
  }

  const handleServerMessage = useCallback((message: WsOutbound) => {
    switch (message.type) {
      case 'session_established': {
        const msg = message as SessionEstablished
        setCurrentSessionId(msg.session_id)
        break
      }

      case 'session_created': {
        const msg = message as SessionCreated
        setCurrentTopicId(msg.session_id)
        setIsLoading(false)
        break
      }

      case 'session_list': {
        const msg = message as SessionList
        // Convert sessions to topics format
        const newTopics = msg.sessions.map((s) => ({
          id: s.session_id,
          session_id: s.session_id,
          title: s.title,
          message_count: Number(s.message_count),
          created_at: s.last_active_at,
          updated_at: s.last_active_at,
        }))
        setTopics(newTopics)
        if (msg.current_session_id) {
          setCurrentTopicId(msg.current_session_id)
        }
        break
      }

      case 'assistant_response': {
        const msg = message as AssistantResponse
        setMessages((prev) => [
          ...prev,
          {
            id: msg.id,
            role: 'assistant',
            content: msg.content,
            timestamp: Date.now(),
            type: 'message',
          },
        ])
        setIsLoading(false)
        break
      }

      case 'tool_call': {
        const msg = message as ToolCall
        setMessages((prev) => [
          ...prev,
          {
            id: msg.id,
            role: 'tool',
            content: msg.content,
            timestamp: Date.now(),
            type: 'tool_call',
            toolName: msg.tool_name,
            arguments: msg.arguments,
          },
        ])
        break
      }

      case 'tool_result': {
        const msg = message as ToolResult
        setMessages((prev) => [
          ...prev,
          {
            id: msg.id,
            role: 'tool',
            content: msg.content,
            timestamp: Date.now(),
            type: 'tool_result',
            toolName: msg.tool_name,
          },
        ])
        break
      }

      case 'tool_pending': {
        const msg = message as ToolPending
        setMessages((prev) => [
          ...prev,
          {
            id: msg.id,
            role: 'tool',
            content: `${msg.content}\n\n${msg.resume_hint}`,
            timestamp: Date.now(),
            type: 'tool_pending',
            toolName: msg.tool_name,
          },
        ])
        break
      }

      case 'error': {
        setMessages((prev) => [
          ...prev,
          {
            id: generateMessageId(),
            role: 'assistant',
            content: `Error: ${message.message}`,
            timestamp: Date.now(),
            type: 'message',
          },
        ])
        setIsLoading(false)
        break
      }
    }
  }, [])

  const handleMessage = useCallback((content: string) => {
    // Add user message to list
    setMessages((prev) => [
      ...prev,
      {
        id: generateMessageId(),
        role: 'user',
        content,
        timestamp: Date.now(),
        type: 'message',
      },
    ])
    setIsLoading(true)
  }, [])

  const handleCommand = useCallback((command: Command) => {
    // Handle local state updates for commands
    switch (command.type) {
      case 'create_session':
        // Optimistically update
        setIsLoading(true)
        break
      case 'list_sessions':
        setIsLoading(true)
        break
      case 'switch_session':
        setCurrentTopicId(command.session_id)
        // Clear messages when switching topic
        setMessages([])
        break
    }
  }, [])

  const clearMessages = useCallback(() => {
    setMessages([])
  }, [])

  return {
    messages,
    currentSessionId,
    currentTopicId,
    topics,
    isLoading,
    handleMessage,
    handleCommand,
    clearMessages,
    handleServerMessage,
  }
}