import { defineStore } from 'pinia'
import { reactive } from 'vue'
import type { EventLogEntry } from '@/types/messages'

interface StreamBuffer {
  text: string
  lastTouched: number // ms since epoch – refreshed on every token
  done?: boolean // set to true after chat_end
  userMessage?: string // the user's question being answered
}

interface EventBuffer {
  liveEvents: Array<{ kind: string; data: any }>
  detailedEventLog: Array<EventLogEntry>
  hasPendingLiveEvent: boolean
  lastTouched: number // ms since epoch – refreshed on every event
}

export const useStreamingStore = defineStore('streaming', {
  state: () => ({
    replyBuffers: reactive<Record<string, StreamBuffer>>({}),
    eventBuffers: reactive<Record<string, EventBuffer>>({}),
  }),

  actions: {
    // Reply buffer methods
    updateReplyBuffer(
      conversationId: string,
      text: string,
      done: boolean = false,
      userMessage?: string
    ) {
      const buf = this.replyBuffers[conversationId] ?? { text: '', lastTouched: Date.now() }
      buf.text = text
      buf.lastTouched = Date.now()
      buf.done = done
      if (userMessage !== undefined) {
        buf.userMessage = userMessage
      }
      this.replyBuffers[conversationId] = buf
    },

    appendToReplyBuffer(conversationId: string, token: string) {
      const buf = this.replyBuffers[conversationId] ?? { text: '', lastTouched: Date.now() }
      buf.text += token
      buf.lastTouched = Date.now()
      this.replyBuffers[conversationId] = buf
    },

    setUserMessage(conversationId: string, userMessage: string) {
      const buf = this.replyBuffers[conversationId] ?? { text: '', lastTouched: Date.now() }
      buf.userMessage = userMessage
      buf.lastTouched = Date.now()
      this.replyBuffers[conversationId] = buf
    },

    markReplyBufferDone(conversationId: string) {
      const buf = this.replyBuffers[conversationId]
      if (buf) {
        buf.done = true
        buf.lastTouched = Date.now()
      }
    },

    getReplyBuffer(conversationId: string): StreamBuffer | undefined {
      return this.replyBuffers[conversationId]
    },

    deleteReplyBuffer(conversationId: string) {
      delete this.replyBuffers[conversationId]
    },

    // Event buffer methods
    addEventToBuffer(conversationId: string, eventKind: string, eventData: any) {
      const eventBuf = this.eventBuffers[conversationId] ?? {
        liveEvents: [],
        detailedEventLog: [],
        hasPendingLiveEvent: false,
        lastTouched: Date.now(),
      }

      eventBuf.liveEvents.push({ kind: eventKind, data: eventData })
      eventBuf.detailedEventLog.push({ eventKind, eventData })
      eventBuf.hasPendingLiveEvent = true
      eventBuf.lastTouched = Date.now()
      this.eventBuffers[conversationId] = eventBuf
    },

    clearPendingEvent(conversationId: string) {
      const eventBuf = this.eventBuffers[conversationId]
      if (eventBuf) {
        eventBuf.hasPendingLiveEvent = false
        eventBuf.lastTouched = Date.now()
      }
    },

    getEventBuffer(conversationId: string): EventBuffer | undefined {
      return this.eventBuffers[conversationId]
    },

    deleteEventBuffer(conversationId: string) {
      delete this.eventBuffers[conversationId]
    },

    // Cleanup methods
    cleanupOldBuffers(currentConversationId?: string) {
      const now = Date.now()
      const fiveMinutes = 300_000 // 5 minutes in milliseconds

      // Clean up old reply buffers
      for (const [cid, buf] of Object.entries(this.replyBuffers)) {
        if (cid !== currentConversationId && buf.done && now - buf.lastTouched > fiveMinutes) {
          delete this.replyBuffers[cid]
        }
      }

      // Clean up old event buffers
      for (const [cid, eventBuf] of Object.entries(this.eventBuffers)) {
        if (cid !== currentConversationId && now - eventBuf.lastTouched > fiveMinutes) {
          delete this.eventBuffers[cid]
        }
      }
    },

    // Clear all buffers (useful for testing or reset)
    clearAllBuffers() {
      Object.keys(this.replyBuffers).forEach((key) => delete this.replyBuffers[key])
      Object.keys(this.eventBuffers).forEach((key) => delete this.eventBuffers[key])
    },
  },

  getters: {
    // Get all active conversation IDs (those with non-done buffers)
    activeConversationIds: (state) => {
      return Object.keys(state.replyBuffers).filter((cid) => !state.replyBuffers[cid]?.done)
    },

    // Get all conversations with pending events
    conversationsWithPendingEvents: (state) => {
      return Object.keys(state.eventBuffers).filter(
        (cid) => state.eventBuffers[cid]?.hasPendingLiveEvent
      )
    },
  },
})
