PicoBot/web/src/hooks/useChat.ts

import { useState, useCallback, useEffect, useRef, useMemo } from 'react'
import type {
  Command,
  ChatMessage,
  Topic,
  WsOutbound,
  AssistantResponse,
  ToolCall,
  ToolResult,
  ToolPending,
  SessionEstablished,
  SessionList,
  TopicList,
  TopicSummary,
  Session,
  TaskMessagesLoaded,
  Attachment,
} from '../types/protocol'

// 简化后的层级状态
interface UseChatReturn {
  // 连接状态
  connectionId: string | null
  isConnected: boolean

  // 简化的层级状态
  session: Session | null
  sessionId: string | null
  chatId: string
  topics: Topic[]
  selectedTopic: string | null

  // 消息
  messages: ChatMessage[]
  isLoading: boolean

  // 是否只读（WebSocket 通道始终可写）
  isReadOnly: boolean

  // 子智能体视图
  subAgentView: SubAgentView | null

  // 方法
  handleMessage: (content: string, attachments?: Attachment[]) => void
  handleCommand: (command: Command) => void
  clearMessages: () => void
  handleServerMessage: (message: WsOutbound) => void

  // Topic 方法
  selectTopic: (topicId: string) => void
  createTopic: (title?: string) => Command
  switchTopic: (topicId: string) => Command

  // 初始化方法
  requestSessionList: () => Command
  requestTopicList: () => Command | null

  // 子智能体导航方法
  enterSubAgentView: (taskId: string, description: string) => Command
  exitSubAgentView: () => void
}

interface SubAgentView {
  taskId: string
  description: string
  subagentType: string
  status: string
  summary?: string
  messages: ChatMessage[]
}

const DEFAULT_CHANNEL = 'websocket'
const DEFAULT_CHAT_ID = 'default'

export function useChat(): UseChatReturn {
  const [messages, setMessages] = useState<ChatMessage[]>([])
  const [isLoading, setIsLoading] = useState(false)

  // 简化的状态管理
  const [connectionId, setConnectionId] = useState<string | null>(null)
  const [session, setSession] = useState<Session | null>(null)
  const [topics, setTopics] = useState<Topic[]>([])
  const [selectedTopic, setSelectedTopic] = useState<string | null>(null)
  const [subAgentView, setSubAgentView] = useState<SubAgentView | null>(null)

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

  // Ref to track subAgentView for use in callbacks
  const subAgentViewRef = useRef<SubAgentView | null>(null)

  const isConnected = useMemo(() => connectionId !== null, [connectionId])
  const sessionId = useMemo(() => session?.id ?? null, [session])
  const chatId = useMemo(() => sessionId ?? DEFAULT_CHAT_ID, [sessionId])

  // Extract subagent_task_id from a message if present
  const getSubagentTaskId = (message: WsOutbound): string | undefined => {
    if (message.type === 'tool_call' || message.type === 'tool_result'
        || message.type === 'tool_pending' || message.type === 'assistant_response') {
      return (message as ToolCall | ToolResult | ToolPending | AssistantResponse).subagent_task_id
    }
    return undefined
  }

  // Convert a server message to ChatMessage (extracted from handleServerMessage logic)
  const serverMessageToChatMessage = (message: WsOutbound): ChatMessage | null => {
    switch (message.type) {
      case 'assistant_response': {
        const msg = message as AssistantResponse
        const role = msg.role === 'user' || msg.role === 'tool' ? msg.role : 'assistant'
        return {
          id: msg.id,
          role: role as ChatMessage['role'],
          content: msg.content,
          timestamp: Date.now(),
          type: 'message',
          attachments: msg.attachments,
          subagentTaskId: msg.subagent_task_id,
        }
      }
      case 'tool_call': {
        const msg = message as ToolCall
        return {
          id: msg.id,
          role: 'tool',
          content: msg.content,
          timestamp: Date.now(),
          type: 'tool_call',
          toolName: msg.tool_name,
          toolCallId: msg.tool_call_id,
          arguments: msg.arguments,
          subagentTaskId: msg.subagent_task_id,
        }
      }
      case 'tool_result': {
        const msg = message as ToolResult
        return {
          id: msg.id,
          role: 'tool',
          content: msg.content,
          timestamp: Date.now(),
          type: 'tool_result',
          toolName: msg.tool_name,
          toolCallId: msg.tool_call_id,
          subagentTaskId: msg.subagent_task_id,
        }
      }
      case 'tool_pending': {
        const msg = message as ToolPending
        return {
          id: msg.id,
          role: 'tool',
          content: `${msg.content}\n\n${msg.resume_hint}`,
          timestamp: Date.now(),
          type: 'tool_pending',
          toolName: msg.tool_name,
          toolCallId: msg.tool_call_id,
          subagentTaskId: msg.subagent_task_id,
        }
      }
      case 'error': {
        return {
          id: generateMessageId(),
          role: 'assistant',
          content: `Error: ${message.message}`,
          timestamp: Date.now(),
          type: 'message',
        }
      }
      default:
        return null
    }
  }

  // Append a server message to the sub-agent view
  const appendToSubAgentViewMessage = (message: WsOutbound) => {
    const chatMsg = serverMessageToChatMessage(message)
    if (chatMsg) {
      setSubAgentView((prev) =>
        prev
          ? { ...prev, messages: [...prev.messages, chatMsg] }
          : prev
      )
    }
  }

  const handleServerMessage = useCallback((message: WsOutbound) => {
    console.log('Received message:', message)

    // Route to sub-agent view if active
    const currentSubAgentView = subAgentViewRef.current
    if (currentSubAgentView) {
      if (message.type === 'task_messages_loaded') {
        const msg = message as TaskMessagesLoaded
        setSubAgentView((prev) =>
          prev
            ? {
                ...prev,
                subagentType: msg.subagent_type,
                status: msg.status,
                summary: msg.summary,
              }
            : prev
        )
        return
      }

      // Route messages to sub-agent view:
      // - Messages without subagent_task_id = loaded history, always accept
      // - Messages with subagent_task_id = live emitter, only accept if matching
      const msgSubagentTaskId = getSubagentTaskId(message)
      if (!msgSubagentTaskId || msgSubagentTaskId === currentSubAgentView.taskId) {
        appendToSubAgentViewMessage(message)
      }
      return
    }

    // In main view, skip sub-agent messages (they belong to sub-agent view).
    // But use the task_id to associate with the running task tool card.
    const msgSubagentTaskId = getSubagentTaskId(message)
    if (msgSubagentTaskId) {
      setMessages((prev) => {
        for (let i = prev.length - 1; i >= 0; i--) {
          if (prev[i].type === 'tool_call' && prev[i].toolName === 'task' && !prev[i].subagentTaskId) {
            const updated = [...prev]
            updated[i] = { ...updated[i], subagentTaskId: msgSubagentTaskId }
            return updated
          }
        }
        return prev
      })
      return
    }

    switch (message.type) {
      case 'session_established': {
        const msg = message as SessionEstablished
        setConnectionId(msg.session_id)
        console.log('Connection established:', msg.session_id)
        break
      }

      case 'session_list': {
        const msg = message as SessionList
        console.log('Session list received:', msg)

        // 自动选择第一个 Session（WebSocket 通道只有一个）
        if (msg.sessions.length > 0) {
          const firstSession = msg.sessions[0]
          setSession({
            id: firstSession.session_id,
            title: firstSession.title,
            channel_name: firstSession.channel_name,
            created_at: Date.now(),
            updated_at: Date.now(),
          })
        }
        setIsLoading(false)
        break
      }

      case 'session_created': {
        // 创建新 Topic 后更新列表
        setIsLoading(false)
        break
      }

      case 'session_loaded': {
        setIsLoading(false)
        break
      }

      case 'topic_list': {
        const msg = message as TopicList
        console.log('Topic list received:', msg)

        // 转换 topics 格式
        const newTopics: Topic[] = msg.topics.map((t: TopicSummary) => ({
          id: t.topic_id,
          session_id: t.session_id,
          title: t.title,
          description: t.description || undefined,
          message_count: Number(t.message_count),
          created_at: t.created_at,
          updated_at: t.last_active_at,
        }))
        setTopics(newTopics)

        // 默认选中第一个 Topic（如果没有选中）
        setIsLoading(false)
        break
      }

      case 'assistant_response': {
        const msg = message as AssistantResponse
        const role = msg.role === 'user' || msg.role === 'tool' ? msg.role : 'assistant'
        setMessages((prev) => [
          ...prev,
          {
            id: msg.id,
            role,
            content: msg.content,
            timestamp: Date.now(),
            type: 'message',
            attachments: msg.attachments,
          },
        ])
        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,
            toolCallId: msg.tool_call_id,
            arguments: msg.arguments,
            subagentTaskId: msg.subagent_task_id,
          },
        ])
        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,
            toolCallId: msg.tool_call_id,
            subagentTaskId: msg.subagent_task_id,
          },
        ])
        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,
            toolCallId: msg.tool_call_id,
          },
        ])
        break
      }

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

      case 'channel_list':
      case 'pong':
        // 忽略这些消息
        break
    }
  }, [])

  const handleMessage = useCallback((content: string, attachments?: Attachment[]) => {
    setMessages((prev) => [
      ...prev,
      {
        id: generateMessageId(),
        role: 'user',
        content,
        timestamp: Date.now(),
        type: 'message',
        attachments: attachments || [],
      },
    ])
    setIsLoading(true)
  }, [])

  const handleCommand = useCallback((command: Command) => {
    switch (command.type) {
      case 'create_session':
      case 'switch_topic':
      case 'load_topic':
      case 'list_sessions':
      case 'list_sessions_by_channel':
      case 'list_topics':
        setIsLoading(true)
        break
    }
  }, [])

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

  // Topic 操作方法
  const selectTopic = useCallback((topicId: string) => {
    setSelectedTopic(topicId)
    setMessages([])
  }, [])

  const createTopic = useCallback((title?: string): Command => {
    return {
      type: 'create_session',
      title: title || `话题 ${new Date().toLocaleString('zh-CN', { month: 'short', day: 'numeric', hour: '2-digit', minute: '2-digit' })}`,
    }
  }, [])

  const switchTopic = useCallback((topicId: string): Command => {
    return {
      type: 'switch_topic',
      topic_id: topicId,
    }
  }, [])

  // 初始化方法
  const requestSessionList = useCallback((): Command => {
    return {
      type: 'list_sessions_by_channel',
      channel_name: DEFAULT_CHANNEL,
      include_archived: false,
    }
  }, [])

  const requestTopicList = useCallback((): Command | null => {
    if (!sessionId) return null
    return {
      type: 'list_topics',
      session_id: sessionId,
    }
  }, [sessionId])

  // Keep ref in sync with state
  useEffect(() => {
    subAgentViewRef.current = subAgentView
  }, [subAgentView])

  const enterSubAgentView = useCallback((taskId: string, description: string): Command => {
    const newView: SubAgentView = {
      taskId,
      description,
      subagentType: '',
      status: 'loading',
      messages: [],
    }
    // Sync ref immediately so WebSocket response routing works correctly
    subAgentViewRef.current = newView
    setSubAgentView(newView)
    return {
      type: 'load_task_messages',
      task_id: taskId,
    }
  }, [])

  const exitSubAgentView = useCallback(() => {
    subAgentViewRef.current = null
    setSubAgentView(null)
  }, [])

  // Memoize messages: when in sub-agent view, return sub-agent messages
  const resolvedMessages = useMemo(() => {
    if (subAgentView) {
      return subAgentView.messages
    }
    return messages
  }, [subAgentView, messages])

  // WebSocket 通道始终可写
  const isReadOnly = false

  return {
    connectionId,
    isConnected,
    session,
    sessionId,
    chatId,
    topics,
    selectedTopic,
    messages: resolvedMessages,
    isLoading,
    isReadOnly,
    subAgentView,
    handleMessage,
    handleCommand,
    clearMessages,
    handleServerMessage,
    selectTopic,
    createTopic,
    switchTopic,
    requestSessionList,
    requestTopicList,
    enterSubAgentView,
    exitSubAgentView,
  }
}