import type { LLMConversation, LLMMessage } from '~/components/ai'

export function useLlmHistories(key: string) {
  const useStore = defineStore(
    `llmConversations/${key}`,
    () => {
      const conversations = ref<LLMConversation[]>([])

      const createConversation = (conversation: LLMConversation) => {
        const index = conversations.value.findIndex(c => c.id === conversation.id)
        if (index >= 0) {
          conversations.value[index] = conversation
        } else {
          conversations.value.unshift(conversation)
        }
      }

      const getConversation = (conversationId: string) => {
        const index = conversations.value.findIndex(c => c.id === conversationId)
        return index >= 0 ? conversations.value[index] : null
      }

      const updateConversation = (conversationId: string, conversation: Partial<LLMConversation>) => {
        const index = conversations.value.findIndex(c => c.id === conversationId)
        if (index >= 0) {
          conversations.value[index] = {
            ...conversations.value[index],
            ...conversation,
          }
        } else {
          createConversation({
            id: conversationId,
            ...conversation,
          } as LLMConversation)
        }
      }

      const removeConversation = (conversationId: string) => {
        const index = conversations.value.findIndex(c => c.id === conversationId)
        if (index >= 0) {
          conversations.value.splice(index, 1)
        }
      }

      const isConversationExist = (conversationId: string) => {
        return conversations.value.some(c => c.id === conversationId)
      }

      const putMessageToConv = (conversationId: string, message: LLMMessage) => {
        const index = conversations.value.findIndex(c => c.id === conversationId)
        if (index >= 0) {
          const conversation = conversations.value[index]
          conversation.messages.push(message)
          conversations.value[index] = conversation
        } else {
          const newConversation: LLMConversation = {
            id: conversationId,
            title: '',
            messages: [message],
          }
          conversations.value.unshift(newConversation)
        }
      }

      const appendChunkToLast = (conversationId: string, chunk: string) => {
        const index = conversations.value.findIndex(c => c.id === conversationId)
        if (index >= 0) {
          const conversation = conversations.value[index]
          const lastMessage = conversation.messages[conversation.messages.length - 1]
          if (lastMessage && lastMessage.role === 'assistant') {
            lastMessage.content += chunk
            conversations.value[index] = conversation
          }
        }
      }

      return {
        conversations,
        createConversation,
        getConversation,
        updateConversation,
        removeConversation,
        putMessageToConv,
        appendChunkToLast,
        isConversationExist
      }
    },
    {
      persist: {
        key: `xshic_llm_histories_${key}`,
        storage: piniaPluginPersistedstate.localStorage(),
        pick: ['conversations'],
      },
    }
  )

  return useStore()
}