ultisuite-client/lib/api/ws.ts

"use client"

import { useEffect } from "react"
import { useQueryClient, type QueryClient } from "@tanstack/react-query"
import type { WsEvent, WsEventType, WsMailPayload } from "./types"
import { ensureAccessToken } from "@/lib/auth/ensure-access-token"
import { useAuthStore } from "./auth-store"
import { useIsDemoApp } from "@/lib/demo/use-is-demo-app"

export type WsEventListener = (evt: WsEvent) => void

function mailPayload(evt: WsEvent): WsMailPayload | null {
  const payload = evt.payload
  if (!payload || typeof payload !== "object") return null
  const messageId = "message_id" in payload ? String(payload.message_id ?? "") : ""
  const accountId = "account_id" in payload ? String(payload.account_id ?? "") : ""
  if (!messageId) return null
  return { message_id: messageId, account_id: accountId }
}

class WebSocketManager {
  private ws: WebSocket | null = null
  private reconnectAttempts = 0
  private maxReconnectDelay = 30_000
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null
  private lastSeq = 0
  private queryClient: QueryClient | null = null
  private listeners = new Set<WsEventListener>()

  init(queryClient: QueryClient) {
    this.queryClient = queryClient
    this.loadLastSeq()
  }

  subscribe(listener: WsEventListener): () => void {
    this.listeners.add(listener)
    return () => this.listeners.delete(listener)
  }

  connect(token: string) {
    if (this.ws?.readyState === WebSocket.OPEN) return

    const baseUrl =
      process.env.NEXT_PUBLIC_WS_URL ??
      (typeof window !== "undefined"
        ? `${window.location.protocol === "https:" ? "wss:" : "ws:"}//${window.location.host}/ws`
        : "")

    const url = `${baseUrl}?token=${encodeURIComponent(token)}&since=${this.lastSeq}`
    this.ws = new WebSocket(url)

    this.ws.onopen = () => {
      this.reconnectAttempts = 0
    }
    this.ws.onmessage = (event) => this.handleMessage(event)
    this.ws.onclose = () => this.scheduleReconnect(token)
    this.ws.onerror = () => {}
  }

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

  private handleMessage(event: MessageEvent) {
    try {
      const evt: WsEvent = JSON.parse(event.data as string)
      if (evt.type === "ws.ping") {
        this.send({ type: "ws.pong", payload: {} })
        return
      }
      if (evt.seq) {
        this.lastSeq = evt.seq
        this.saveLastSeq()
      }
      this.handleEvent(evt)
      for (const listener of this.listeners) {
        listener(evt)
      }
    } catch {}
  }

  private send(evt: Pick<WsEvent, "type" | "payload">) {
    if (this.ws?.readyState !== WebSocket.OPEN) return
    this.ws.send(JSON.stringify(evt))
  }

  private handleEvent(evt: WsEvent) {
    if (!this.queryClient) return

    const mail = mailPayload(evt)

    switch (evt.type as WsEventType) {
      case "mail.created":
        this.queryClient.invalidateQueries({ queryKey: ["messages"] })
        break
      case "mail.updated":
        this.queryClient.invalidateQueries({ queryKey: ["messages"] })
        if (mail?.message_id) {
          this.queryClient.invalidateQueries({
            queryKey: ["message", mail.message_id],
          })
        }
        break
      case "mail.deleted":
        this.queryClient.invalidateQueries({ queryKey: ["messages"] })
        if (mail?.message_id) {
          this.queryClient.removeQueries({
            queryKey: ["message", mail.message_id],
          })
        }
        break
      case "outbox.updated":
        this.queryClient.invalidateQueries({ queryKey: ["outbox"] })
        break
      case "contact.updated":
        this.queryClient.invalidateQueries({ queryKey: ["contacts"] })
        break
    }
  }

  private scheduleReconnect(token: string) {
    const delay = Math.min(
      1000 * 2 ** this.reconnectAttempts,
      this.maxReconnectDelay
    )
    this.reconnectAttempts++
    this.reconnectTimer = setTimeout(() => this.connect(token), delay)
  }

  private loadLastSeq() {
    if (typeof window === "undefined") return
    const stored = localStorage.getItem("ultimail-ws-seq")
    if (stored) this.lastSeq = parseInt(stored, 10) || 0
  }

  private saveLastSeq() {
    if (typeof window === "undefined") return
    localStorage.setItem("ultimail-ws-seq", String(this.lastSeq))
  }
}

export const wsManager = new WebSocketManager()

export function useWebSocket() {
  const queryClient = useQueryClient()
  const accessToken = useAuthStore((s) => s.accessToken)
  const isDemoApp = useIsDemoApp()

  useEffect(() => {
    if (isDemoApp) return
    wsManager.init(queryClient)
  }, [queryClient, isDemoApp])

  useEffect(() => {
    if (isDemoApp) return
    let cancelled = false

    void (async () => {
      const token = accessToken ? await ensureAccessToken() : null
      if (cancelled) return
      if (token) {
        wsManager.connect(token)
      } else {
        wsManager.disconnect()
      }
    })()

    return () => {
      cancelled = true
      wsManager.disconnect()
    }
  }, [accessToken, isDemoApp])
}

export function useWsEventListener(listener: WsEventListener) {
  useEffect(() => wsManager.subscribe(listener), [listener])
}