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

    // A 401 from the auth endpoints themselves is an answer, not an expired
    // session. Logging in with the wrong password returns 401, and treating
    // that as "your token lapsed" sent it round the refresh path, found no
    // refresh token to use, logged out and reloaded the page - which wiped the
    // error the form had just rendered. The red text was there; the reload
    // took it away before it could be read. Let these fall through so the
    // caller can show what the server actually said.
    const url = originalRequest?.url || ''
    const isAuthAttempt = url.includes('/auth/login') || url.includes('/auth/refresh')

    if (error.response?.status === 401 && !originalRequest._retry && !isAuthAttempt) {
      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>()

  // This connection's own id, told to us when the server accepts the token.
  //
  // Kept here rather than in whichever component happened to be listening:
  // auth_success arrives once, and a page navigated to later has missed it.
  // Something that needs to recognise its own messages - a node lock, say -
  // would then never know which ones were its own.
  clientId: string | null = null
  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.clientId = message.clientId || null
          if (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()