import { useCallback, useMemo, useState } from "react"
import type { UiChatMessage } from "../../lib/chat-utils"

export type MessagesBySession = Record<string, UiChatMessage[]>

export function useSessionMessageStore(visibleSessionId: string) {
  const [messagesBySession, setMessagesBySession] = useState<MessagesBySession>({})
  const messages = useMemo(
    () => (visibleSessionId ? messagesBySession[visibleSessionId] ?? [] : []),
    [messagesBySession, visibleSessionId]
  )

  const updateSessionMessages = useCallback((sessionId: string, updater: (current: UiChatMessage[]) => UiChatMessage[]) => {
    if (!sessionId) {
      return
    }

    setMessagesBySession((current) => {
      const currentMessages = current[sessionId] ?? []
      const nextMessages = updater(currentMessages)
      if (nextMessages === currentMessages) {
        return current
      }
      if (!nextMessages.length) {
        if (!(sessionId in current)) {
          return current
        }
        const { [sessionId]: _removed, ...rest } = current
        return rest
      }
      return {
        ...current,
        [sessionId]: nextMessages
      }
    })
  }, [])

  const clearAllSessionMessages = useCallback(() => {
    setMessagesBySession({})
  }, [])

  const moveSessionMessages = useCallback((sourceSessionId: string, targetSessionId: string, messageIds: string[]) => {
    if (!sourceSessionId || !targetSessionId || sourceSessionId === targetSessionId || messageIds.length === 0) {
      return
    }

    setMessagesBySession((current) => {
      const sourceMessages = current[sourceSessionId] ?? []
      const targetMessages = current[targetSessionId] ?? []
      const movedMessages = sourceMessages.filter((message) => messageIds.includes(message.id))
      if (!movedMessages.length) {
        return current
      }

      const nextSourceMessages = sourceMessages.filter((message) => !messageIds.includes(message.id))
      const targetMessageIds = new Set(targetMessages.map((message) => message.id))
      const nextMovedMessages = movedMessages.filter((message) => !targetMessageIds.has(message.id))

      return {
        ...current,
        [sourceSessionId]: nextSourceMessages,
        [targetSessionId]: nextMovedMessages.length > 0 ? [...targetMessages, ...nextMovedMessages] : targetMessages
      }
    })
  }, [])

  const updateMessageById = useCallback((messageId: string, updater: (message: UiChatMessage) => UiChatMessage) => {
    setMessagesBySession((current) => {
      for (const [sessionId, sessionMessages] of Object.entries(current)) {
        const messageIndex = sessionMessages.findIndex((message) => message.id === messageId)
        if (messageIndex < 0) {
          continue
        }

        const nextMessages = [...sessionMessages]
        nextMessages[messageIndex] = updater(sessionMessages[messageIndex])
        return {
          ...current,
          [sessionId]: nextMessages
        }
      }

      return current
    })
  }, [])

  return {
    messagesBySession,
    messages,
    updateSessionMessages,
    clearAllSessionMessages,
    moveSessionMessages,
    updateMessageById
  }
}
