loyalty-agent-service/src/services/StompService.ts

import { Client } from '@stomp/stompjs'
import { useConnectionStore } from '../store/useConnectionStore'


export interface ChatEvent {
  messageId?: string
  type: 'TOKEN' | 'TOOL_STATUS' | 'DONE' | 'ERROR'
  content?: string
  tool?: string
  status?: string
}

interface QueuedMessage {
  destination: string
  body: string
}

class StompService {
  private client: Client | null = null
  private reconnectAttempt = 0
  private reconnectTimeoutId: any = null
  private messageQueue: QueuedMessage[] = []
  private messageListeners: Array<(event: ChatEvent) => void> = []
  private maxQueueSize = 20

  private getWsUrl(): string {
    const host = window.location.hostname || 'localhost'
    const port = window.location.port ? window.location.port : '9332'
    // If dev server port 9333 or similar, connect to 9332 if port is 9333, else same port
    const targetPort = port === '9333' || port === '5173' ? '9332' : port
    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
    return `${protocol}//${host}:${targetPort}/ws`
  }

  public connect() {
    if (this.client && this.client.active) {
      return
    }

    useConnectionStore.getState().setStatus('connecting')

    this.client = new Client({
      brokerURL: this.getWsUrl(),
      reconnectDelay: 0, // Manual reconnect management for exponential backoff + jitter
      heartbeatIncoming: 10000,
      heartbeatOutgoing: 10000,
      onConnect: () => {
        console.log('[StompService] Connected to WebSocket')
        this.reconnectAttempt = 0
        useConnectionStore.getState().setStatus('connected')

        // Subscribe to user chat events
        this.client?.subscribe('/user/queue/chat-events', (message) => {
          try {
            const event: ChatEvent = JSON.parse(message.body)
            this.notifyMessageListeners(event)
          } catch (e) {
            console.error('[StompService] Failed to parse message body:', e)
          }
        })

        // Flush offline queue
        this.flushQueue()
      },
      onWebSocketClose: () => {
        console.warn('[StompService] WebSocket connection closed')
        useConnectionStore.getState().setStatus('disconnected')
        this.scheduleReconnect()
      },
      onStompError: (frame) => {
        console.error('[StompService] STOMP error frame received:', frame.headers['message'])
        useConnectionStore.getState().setStatus('error')
        this.scheduleReconnect()
      },
      onWebSocketError: (event) => {
        console.error('[StompService] WebSocket error:', event)
        useConnectionStore.getState().setStatus('error')
      }
    })

    this.client.activate()
  }

  public disconnect() {
    if (this.reconnectTimeoutId) {
      clearTimeout(this.reconnectTimeoutId)
      this.reconnectTimeoutId = null
    }
    if (this.client) {
      this.client.deactivate()
      this.client = null
    }
    useConnectionStore.getState().setStatus('disconnected')
  }

  private scheduleReconnect() {
    if (this.reconnectTimeoutId) return

    // Base exponential backoff: 1s, 2s, 4s, 8s, 16s, capped at 30s
    const baseDelay = Math.min(1000 * Math.pow(2, this.reconnectAttempt), 30000)
    // Jitter: ±50% (0.5 to 1.5 multiplier)
    const jitter = 0.5 + Math.random()
    const delay = Math.round(baseDelay * jitter)

    this.reconnectAttempt++
    console.log(`[StompService] Scheduling reconnect attempt #${this.reconnectAttempt} in ${delay}ms`)

    this.reconnectTimeoutId = setTimeout(() => {
      this.reconnectTimeoutId = null
      if (this.client) {
        this.client.activate()
      } else {
        this.connect()
      }
    }, delay)
  }

  public publish(destination: string, payload: object): boolean {
    const body = JSON.stringify(payload)
    const isConnected = this.client && this.client.connected

    if (isConnected) {
      this.client!.publish({ destination, body })
      return true
    } else {
      if (this.messageQueue.length >= this.maxQueueSize) {
        console.warn('[StompService] Queue limit reached (20). Message dropped.')
        return false
      }
      this.messageQueue.push({ destination, body })
      useConnectionStore.getState().setQueuedCount(this.messageQueue.length)
      console.log(`[StompService] Message queued offline (${this.messageQueue.length}/${this.maxQueueSize})`)
      
      // Ensure we are trying to reconnect
      this.connect()
      return true
    }
  }

  private flushQueue() {
    if (!this.client || !this.client.connected) return

    while (this.messageQueue.length > 0) {
      const msg = this.messageQueue.shift()
      if (msg) {
        this.client.publish({ destination: msg.destination, body: msg.body })
      }
    }
    useConnectionStore.getState().setQueuedCount(0)
    console.log('[StompService] Offline message queue flushed')
  }

  public onMessage(callback: (event: ChatEvent) => void): () => void {
    this.messageListeners.push(callback)
    return () => {
      this.messageListeners = this.messageListeners.filter((cb) => cb !== callback)
    }
  }

  private notifyMessageListeners(event: ChatEvent) {
    this.messageListeners.forEach((listener) => {
      try {
        listener(event)
      } catch (e) {
        console.error('[StompService] Error in message listener:', e)
      }
    })
  }

  public getQueueLength(): number {
    return this.messageQueue.length
  }
}

export const stompService = new StompService()