import { create } from 'zustand'
import { stompService } from '../services/StompService'
import type { ChatEvent } from '../services/StompService'


export type Role = 'user' | 'agent'

export interface Message {
  id: string
  role: Role
  content: string
  isStreaming?: boolean
  toolStatus?: string
  createdAt?: string
}

interface ChatState {
  conversationId: string
  messages: Message[]
  activeAgentMessageId: string | null
  setConversationId: (id: string) => void
  setMessages: (messages: Message[]) => void
  initStomp: () => void
  sendMessage: (content: string, activeCampaignId?: string) => void
  addMessage: (message: Message) => void
  updateMessage: (id: string, updater: (msg: Message) => Message) => void
}

let isStompInitialized = false

export const useChatStore = create<ChatState>((set, get) => ({
  conversationId: '',
  messages: [],
  activeAgentMessageId: null,

  setConversationId: (id: string) => set({ conversationId: id }),

  setMessages: (messages: Message[]) => set({ messages }),

  initStomp: () => {
    if (isStompInitialized) return
    isStompInitialized = true

    stompService.connect()
    stompService.onMessage((event: ChatEvent) => {
      const msgId = event.messageId
      if (!msgId) return

      if (event.type === 'TOKEN') {
        get().updateMessage(msgId, (msg) => ({
          ...msg,
          content: msg.content + (event.content || ''),
          isStreaming: true
        }))
      } else if (event.type === 'TOOL_STATUS') {
        get().updateMessage(msgId, (msg) => ({
          ...msg,
          toolStatus: event.status === 'running' ? event.tool || 'running' : ''
        }))
      } else if (event.type === 'DONE') {
        get().updateMessage(msgId, (msg) => ({
          ...msg,
          isStreaming: false,
          toolStatus: undefined
        }))
      } else if (event.type === 'ERROR') {
        get().updateMessage(msgId, (msg) => ({
          ...msg,
          content: msg.content + '\n[ERROR] ' + (event.content || ''),
          isStreaming: false,
          toolStatus: undefined
        }))
      }
    })
  },

  sendMessage: (content: string, activeCampaignId?: string) => {
    const userMessage: Message = {
      id: crypto.randomUUID(),
      role: 'user',
      content,
    }
    
    get().addMessage(userMessage)

    const agentMessageId = crypto.randomUUID()
    set({ activeAgentMessageId: agentMessageId })
    get().addMessage({
      id: agentMessageId,
      role: 'agent',
      content: '',
      isStreaming: true
    })

    const payload = {
      conversationId: get().conversationId,
      messageId: agentMessageId,
      prompt: content,
      activeCampaignId
    }

    stompService.publish('/app/chat', payload)
  },

  addMessage: (message: Message) => {
    set((state) => ({ messages: [...state.messages, message] }))
  },

  updateMessage: (id: string, updater: (msg: Message) => Message) => {
    set((state) => ({
      messages: state.messages.map((m) => (m.id === id ? updater(m) : m)),
    }))
  },
}))
