smartbotic/webui/src/api/client.ts

import axios, { AxiosError } from 'axios'
import { useAuthStore } from '../stores/authStore'

export const api = axios.create({
  baseURL: '/api/v1',
  headers: {
    'Content-Type': 'application/json',
  },
  timeout: 30000, // 30 second timeout
})

// Track consecutive network failures
let networkFailureCount = 0
const MAX_NETWORK_FAILURES = 3

// Request interceptor - add auth token
api.interceptors.request.use((config) => {
  const token = useAuthStore.getState().accessToken
  if (token) {
    config.headers.Authorization = `Bearer ${token}`
  }
  return config
})

// Response interceptor - handle token refresh and network errors
api.interceptors.response.use(
  (response) => {
    // Reset failure count on success
    networkFailureCount = 0
    return response
  },
  async (error: AxiosError) => {
    const originalRequest = error.config as any

    // Handle network errors (server unreachable, timeout, etc.)
    if (!error.response) {
      networkFailureCount++
      console.error(`Network error (${networkFailureCount}/${MAX_NETWORK_FAILURES}):`, error.message)

      // Auto-logout after consecutive failures
      if (networkFailureCount >= MAX_NETWORK_FAILURES && useAuthStore.getState().accessToken) {
        console.warn('Server appears to be down, logging out...')
        useAuthStore.getState().logout()
        window.location.href = '/login?reason=server_unavailable'
        return Promise.reject(new Error('Server unavailable'))
      }

      return Promise.reject(error)
    }

    // Reset failure count when we get a response (even errors)
    networkFailureCount = 0

    if (error.response?.status === 401 && !originalRequest._retry) {
      originalRequest._retry = true

      try {
        await useAuthStore.getState().refreshTokens()
        const token = useAuthStore.getState().accessToken
        originalRequest.headers.Authorization = `Bearer ${token}`
        return api(originalRequest)
      } catch (refreshError) {
        useAuthStore.getState().logout()
        window.location.href = '/login'
        return Promise.reject(refreshError)
      }
    }

    return Promise.reject(error)
  }
)

// WebSocket connection
export class WebSocketClient {
  private ws: WebSocket | null = null
  private reconnectTimeout: number | null = null
  private messageHandlers: Map<string, Set<(data: any) => void>> = new Map()
  // What this client wants to be subscribed to, as opposed to what it has
  // managed to say so far. send() drops anything queued before the socket is
  // open, so a page that subscribes while connecting - which is every page,
  // on a cold load - was silently never subscribed to anything.
  private desiredChannels = new Set<string>()
  private reconnectAttempts: number = 0
  private maxReconnectAttempts: number = 5
  private baseReconnectDelay: number = 1000
  private maxReconnectDelay: number = 30000
  private reconnectCallbacks: Set<() => void> = new Set()
  private wasConnected: boolean = false

  connect() {
    // Don't try to connect if not logged in
    if (!useAuthStore.getState().accessToken) {
      return
    }

    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'

    // In development (Vite), use /ws path which is proxied to the backend
    // In production, WebSocket server runs on HTTP port + 1
    let wsUrl: string
    if (import.meta.env.DEV) {
      // Use Vite proxy - connects to same host/port with /ws path
      wsUrl = `${protocol}//${window.location.host}/ws`
    } else {
      // Production: WebSocket server on HTTP port + 1
      const httpPort = parseInt(window.location.port) || (window.location.protocol === 'https:' ? 443 : 80)
      const wsPort = httpPort + 1
      wsUrl = `${protocol}//${window.location.hostname}:${wsPort}`
    }

    try {
      this.ws = new WebSocket(wsUrl)
    } catch (e) {
      console.error('Failed to create WebSocket:', e)
      this.handleReconnect()
      return
    }

    this.ws.onopen = () => {
      console.log('WebSocket connected')
      const isReconnect = this.wasConnected
      this.reconnectAttempts = 0 // Reset on successful connection
      this.wasConnected = true

      // Authenticate
      const token = useAuthStore.getState().accessToken
      if (token) {
        this.send({ type: 'auth', token: `Bearer ${token}` })
      }

      // Notify listeners about reconnection so they can reset state
      if (isReconnect) {
        console.log('WebSocket reconnected - notifying listeners to reset state')
        this.reconnectCallbacks.forEach(callback => {
          try {
            callback()
          } catch (e) {
            console.error('Error in reconnect callback:', e)
          }
        })
      }
    }

    this.ws.onmessage = (event) => {
      try {
        const message = JSON.parse(event.data)

        // Subscriptions are replayed here rather than in onopen because the
        // server refuses a subscribe until it has accepted the token.
        if (message.type === 'auth_success' && this.desiredChannels.size > 0) {
          this.send({ type: 'subscribe', channels: [...this.desiredChannels] })
        }

        const channel = message.channel || message.type
        const data = message.data || message

        // Include channel in data for wildcard handlers
        const dataWithChannel = { ...data, _channel: channel }

        const handlers = this.messageHandlers.get(channel)
        if (handlers) {
          handlers.forEach((handler) => handler(data))
        }

        // Wildcard handlers - pass channel info so handlers know event type
        this.messageHandlers.forEach((handlers, pattern) => {
          if (pattern.endsWith('.*') && channel.startsWith(pattern.slice(0, -1))) {
            handlers.forEach((handler) => handler(dataWithChannel))
          }
        })
      } catch (e) {
        console.error('WebSocket message parse error:', e)
      }
    }

    this.ws.onclose = () => {
      console.log('WebSocket disconnected')
      this.handleReconnect()
    }

    this.ws.onerror = (error) => {
      console.error('WebSocket error:', error)
    }
  }

  private handleReconnect() {
    // Don't reconnect if not logged in
    if (!useAuthStore.getState().accessToken) {
      return
    }

    this.reconnectAttempts++

    if (this.reconnectAttempts > this.maxReconnectAttempts) {
      console.warn('Max WebSocket reconnect attempts reached, server appears down')
      useAuthStore.getState().logout()
      window.location.href = '/login?reason=server_unavailable'
      return
    }

    // Exponential backoff with jitter
    const delay = Math.min(
      this.baseReconnectDelay * Math.pow(2, this.reconnectAttempts - 1) + Math.random() * 1000,
      this.maxReconnectDelay
    )

    console.log(`WebSocket reconnecting in ${Math.round(delay)}ms (attempt ${this.reconnectAttempts}/${this.maxReconnectAttempts})`)
    this.reconnectTimeout = window.setTimeout(() => this.connect(), delay)
  }

  disconnect() {
    if (this.reconnectTimeout) {
      clearTimeout(this.reconnectTimeout)
    }
    if (this.ws) {
      this.ws.close()
      this.ws = null
    }
  }

  send(message: object) {
    if (this.ws?.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify(message))
    }
  }

  subscribe(channels: string[]) {
    channels.forEach((c) => this.desiredChannels.add(c))
    this.send({ type: 'subscribe', channels })
  }

  unsubscribe(channels: string[]) {
    channels.forEach((c) => this.desiredChannels.delete(c))
    this.send({ type: 'unsubscribe', channels })
  }

  on(channel: string, handler: (data: any) => void) {
    if (!this.messageHandlers.has(channel)) {
      this.messageHandlers.set(channel, new Set())
    }
    this.messageHandlers.get(channel)!.add(handler)
  }

  off(channel: string, handler: (data: any) => void) {
    this.messageHandlers.get(channel)?.delete(handler)
  }

  // Register a callback to be called when WebSocket reconnects
  // Useful for components to reset their execution state
  onReconnect(callback: () => void) {
    this.reconnectCallbacks.add(callback)
  }

  offReconnect(callback: () => void) {
    this.reconnectCallbacks.delete(callback)
  }
}

export const wsClient = new WebSocketClient()