juwan-frontend/lib/hooks/use-chat-socket.ts

"use client"

import { useAuthStore } from "@/store/auth"
import { useCallback, useEffect, useRef, useState } from "react"

export type ChatSocketMessage = {
  id: string
  sessionId: string
  senderId: string
  type: "text" | "image" | "system"
  content: string
  createdAt: string
}

type WsEnvelope = {
  type: string
  sessionId?: number
  senderId?: number
  content?: string
  msgType?: string
  data?: unknown
}

type ChatSocketState = {
  connected: boolean
  sessionId: string | null
  messages: ChatSocketMessage[]
  error: string | null
  createDM: (targetUserId: string) => void
  joinSession: (nextSessionId: string) => void
  leaveSession: (nextSessionId: string) => void
  sendTextMessage: (content: string) => void
  sendImageMessage: (imageUrl: string) => void
  requestHistory: (nextSessionId: string) => void
}

function getChatSocketUrl() {
  const configured = process.env.NEXT_PUBLIC_BACKEND_URL
  const origin =
    configured ||
    (process.env.NODE_ENV === "development" ? "http://localhost:18080" : window.location.origin)
  const url = new URL(origin)
  url.protocol = url.protocol === "https:" ? "wss:" : "ws:"
  url.pathname = "/ws/chat"
  url.search = ""
  return url.toString()
}

function unixToIso(value: unknown) {
  if (typeof value !== "number" || value <= 0) return new Date().toISOString()
  return new Date(value * 1000).toISOString()
}

function normalizeMessage(value: unknown): ChatSocketMessage | null {
  if (typeof value !== "object" || value === null) return null
  const item = value as {
    id?: unknown
    sessionId?: unknown
    senderId?: unknown
    type?: unknown
    content?: unknown
    createdAt?: unknown
  }
  if (typeof item.content !== "string") return null

  const type = item.type === "image" || item.type === "system" ? item.type : "text"
  return {
    id: String(item.id ?? `${item.sessionId ?? ""}-${item.senderId ?? ""}-${item.createdAt ?? ""}`),
    sessionId: String(item.sessionId ?? ""),
    senderId: String(item.senderId ?? ""),
    type,
    content: item.content,
    createdAt: unixToIso(item.createdAt),
  }
}

function liveMessage(msg: WsEnvelope): ChatSocketMessage | null {
  if (msg.type !== "message" || !msg.content || msg.sessionId === undefined) return null
  const data = typeof msg.data === "object" && msg.data !== null ? msg.data : {}
  const messageId = "messageId" in data ? (data as { messageId?: unknown }).messageId : undefined
  const type = msg.msgType === "image" ? "image" : "text"

  return {
    id: String(messageId ?? `${msg.sessionId}-${msg.senderId ?? ""}-${Date.now()}`),
    sessionId: String(msg.sessionId),
    senderId: String(msg.senderId ?? ""),
    type,
    content: msg.content,
    createdAt: new Date().toISOString(),
  }
}

function jsonWithInt64(type: string, fields: Record<string, string | undefined>) {
  const parts = [`"type":${JSON.stringify(type)}`]
  for (const [key, value] of Object.entries(fields)) {
    if (!value) continue
    parts.push(`"${key}":${value}`)
  }
  return `{${parts.join(",")}}`
}

function jsonMessage(sessionId: string, content: string, msgType: "text" | "image") {
  const payload = JSON.stringify({ type: "message", content, msgType })
  return payload.replace(/^\{/, `{"sessionId":${sessionId},`)
}

export function useChatSocket(): ChatSocketState {
  const wsRef = useRef<WebSocket | null>(null)
  const [connected, setConnected] = useState(false)
  const [sessionId, setSessionId] = useState<string | null>(null)
  const [messages, setMessages] = useState<ChatSocketMessage[]>([])
  const [error, setError] = useState<string | null>(null)
  const userId = useAuthStore((state) => state.user?.id)

  useEffect(() => {
    if (!userId) return

    const ws = new WebSocket(getChatSocketUrl())
    wsRef.current = ws

    ws.onopen = () => {
      if (wsRef.current !== ws) return
      setConnected(true)
      setError(null)
    }

    ws.onmessage = (event) => {
      const msg = JSON.parse(event.data) as WsEnvelope

      if (msg.type === "dm_created" || msg.type === "group_created" || msg.type === "joined") {
        if (msg.sessionId !== undefined) setSessionId(String(msg.sessionId))
        if (msg.sessionId !== undefined)
          ws.send(jsonWithInt64("history", { sessionId: String(msg.sessionId) }))
        return
      }

      if (msg.type === "history") {
        const nextMessages = Array.isArray(msg.data)
          ? msg.data
              .map(normalizeMessage)
              .filter((item): item is ChatSocketMessage => item !== null)
          : []
        setMessages(nextMessages)
        return
      }

      if (msg.type === "error") {
        setError(msg.content ?? "聊天服务错误")
        return
      }

      const nextMessage = liveMessage(msg)
      if (nextMessage) setMessages((prev) => [...prev, nextMessage])
    }

    ws.onclose = () => {
      if (wsRef.current !== ws) return
      setConnected(false)
      wsRef.current = null
    }

    ws.onerror = () => {
      if (wsRef.current !== ws) return
      setError("聊天连接失败")
      ws.close()
    }

    return () => {
      if (wsRef.current === ws) wsRef.current = null
      ws.close()
    }
  }, [userId])

  const sendRaw = useCallback((payload: string) => {
    const ws = wsRef.current
    if (ws?.readyState === WebSocket.OPEN) ws.send(payload)
  }, [])

  const createDM = useCallback(
    (targetUserId: string) => sendRaw(jsonWithInt64("create_dm", { targetId: targetUserId })),
    [sendRaw],
  )

  const joinSession = useCallback(
    (nextSessionId: string) => sendRaw(jsonWithInt64("join", { sessionId: nextSessionId })),
    [sendRaw],
  )

  const leaveSession = useCallback(
    (nextSessionId: string) => sendRaw(jsonWithInt64("leave", { sessionId: nextSessionId })),
    [sendRaw],
  )

  const requestHistory = useCallback(
    (nextSessionId: string) => sendRaw(jsonWithInt64("history", { sessionId: nextSessionId })),
    [sendRaw],
  )

  const sendTextMessage = useCallback(
    (content: string) => {
      if (!sessionId) return
      sendRaw(jsonMessage(sessionId, content, "text"))
    },
    [sendRaw, sessionId],
  )

  const sendImageMessage = useCallback(
    (imageUrl: string) => {
      if (!sessionId) return
      sendRaw(jsonMessage(sessionId, imageUrl, "image"))
    },
    [sendRaw, sessionId],
  )

  return {
    connected,
    sessionId,
    messages,
    error,
    createDM,
    joinSession,
    leaveSession,
    sendTextMessage,
    sendImageMessage,
    requestHistory,
  }
}