NextEdu/src/modules/ai/hooks/use-ai-chat-stream.ts

"use client"

import { useState, useCallback, useRef, useEffect } from "react"
import { useTranslations } from "next-intl"
import type { AiChatMessage } from "../types"
import {
  consumeSseStream,
  getStreamErrorKey,
  extractErrorMessage,
  removeTrailingEmptyAssistant,
  appendTokenToLastAssistant,
} from "./stream-utils"

const HISTORY_STORAGE_KEY = "ai-chat-history"
const MAX_HISTORY = 20

/**
 * AI 流式聊天 Hook
 *
 * 通过 SSE 端点消费流式 AI 回复。
 * 支持：逐 token 渲染、停止生成（AbortController）、错误处理、历史持久化恢复。
 */
type UseAiChatStreamReturn = {
  messages: AiChatMessage[]
  streaming: boolean
  error: string | null
  send: (messages: AiChatMessage[], options?: { systemPrompt?: string; providerId?: string }) => Promise<void>
  stop: () => void
  clear: () => void
}

function loadHistory(): AiChatMessage[] {
  if (typeof window === "undefined") return []
  try {
    const raw = localStorage.getItem(HISTORY_STORAGE_KEY)
    if (!raw) return []
    const parsed = JSON.parse(raw) as AiChatMessage[]
    if (!Array.isArray(parsed)) return []
    // 过滤掉空的 assistant 消息
    return parsed.filter((m) => m && m.role && typeof m.content === "string")
  } catch {
    return []
  }
}

export function useAiChatStream(): UseAiChatStreamReturn {
  const t = useTranslations("ai")
  // 懒初始化：从 localStorage 恢复历史
  const [messages, setMessages] = useState<AiChatMessage[]>(() => loadHistory())
  const [streaming, setStreaming] = useState(false)
  const [error, setError] = useState<string | null>(null)
  const abortControllerRef = useRef<AbortController | null>(null)

  // 持久化（防抖）
  const persistTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null)
  useEffect(() => {
    if (streaming) return
    if (persistTimerRef.current) clearTimeout(persistTimerRef.current)
    persistTimerRef.current = setTimeout(() => {
      try {
        if (messages.length === 0) {
          localStorage.removeItem(HISTORY_STORAGE_KEY)
        } else {
          localStorage.setItem(HISTORY_STORAGE_KEY, JSON.stringify(messages.slice(-MAX_HISTORY)))
        }
      } catch {
        // ignore
      }
    }, 500)
    return () => {
      if (persistTimerRef.current) clearTimeout(persistTimerRef.current)
    }
  }, [messages, streaming])

  const send = useCallback(
    async (
      inputMessages: AiChatMessage[],
      options?: { systemPrompt?: string; providerId?: string }
    ): Promise<void> => {
      if (streaming) return
      setStreaming(true)
      setError(null)

      const userMessage = inputMessages[inputMessages.length - 1]
      if (userMessage && userMessage.role === "user") {
        setMessages((prev) => [...prev, userMessage])
      }
      setMessages((prev) => [...prev, { role: "assistant", content: "" }])

      const controller = new AbortController()
      abortControllerRef.current = controller

      try {
        const response = await fetch("/api/ai/chat/stream", {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({
            messages: inputMessages,
            systemPrompt: options?.systemPrompt,
            providerId: options?.providerId,
          }),
          signal: controller.signal,
        })

        if (!response.ok) {
          const fallback = t(getStreamErrorKey(response.status))
          const errorMessage = await extractErrorMessage(response, fallback)
          setError(errorMessage)
          setMessages((prev) => removeTrailingEmptyAssistant(prev))
          return
        }

        await consumeSseStream(response, (event) => {
          if (event.type === "token") {
            setMessages((prev) => appendTokenToLastAssistant(prev, event.content))
          } else if (event.type === "error") {
            setError(event.message)
            setMessages((prev) => removeTrailingEmptyAssistant(prev))
          } else if (event.type === "filtered") {
            setError(t("safety.contentFiltered"))
          } else if (event.type === "socratic_warning") {
            // 苏格拉底式辅导警告：不阻断，仅提示
            // 可通过 toast 或 UI 标记展示，此处暂存 error 供组件判断
            setError(event.message)
          }
        })
      } catch (err) {
        if (!(err instanceof DOMException && err.name === "AbortError")) {
          setError(err instanceof Error ? err.message : String(err))
        }
      } finally {
        setStreaming(false)
        abortControllerRef.current = null
        setMessages((prev) => removeTrailingEmptyAssistant(prev))
      }
    },
    [streaming, t]
  )

  const stop = useCallback((): void => {
    abortControllerRef.current?.abort()
  }, [])

  const clear = useCallback((): void => {
    setMessages([])
    setError(null)
    try {
      localStorage.removeItem(HISTORY_STORAGE_KEY)
    } catch {
      // ignore
    }
  }, [])

  return { messages, streaming, error, send, stop, clear }
}