Real-time WebSockets

import { redis, createSubscriber } from "../lib/redis.js";
import { logger } from "../lib/logger.js";

export interface WsClient {
  ws: any;
  rooms: Set<string>;
  userId?: string;
  sessionId?: string;
}

export class WebSocketManager {
  private clients = new Map<string, WsClient>();
  private rooms = new Map<string, Set<string>>();
  private subscriber;

  constructor() {
    this.subscriber = createSubscriber();
    this.setupSubscriber();
  }

  private setupSubscriber() {
    this.subscriber.on("message", (channel: string, message: string) => {
      this.broadcastToRoom(channel, message);
    });
  }

  getClient(id: string): WsClient | undefined {
    return this.clients.get(id);
  }

  addClient(id: string, ws: any, userId?: string, sessionId?: string) {
    this.clients.set(id, { ws, rooms: new Set(), userId, sessionId });
  }

  removeClient(id: string) {
    const client = this.clients.get(id);
    if (client) {
      for (const room of client.rooms) {
        this.leaveRoom(id, room);
      }
      this.clients.delete(id);
    }
  }

  async joinRoom(clientId: string, room: string) {
    const client = this.clients.get(clientId);
    if (!client) return;

    client.rooms.add(room);

    if (!this.rooms.has(room)) {
      this.rooms.set(room, new Set());
      await this.subscriber.subscribe(room);
    }
    this.rooms.get(room)!.add(clientId);
  }

  async leaveRoom(clientId: string, room: string) {
    const client = this.clients.get(clientId);
    if (client) client.rooms.delete(room);

    const roomClients = this.rooms.get(room);
    if (roomClients) {
      roomClients.delete(clientId);
      if (roomClients.size === 0) {
        this.rooms.delete(room);
        await this.subscriber.unsubscribe(room);
      }
    }
  }

  private broadcastToRoom(room: string, message: string) {
    const roomClients = this.rooms.get(room);
    if (!roomClients) return;

    for (const clientId of roomClients) {
      const client = this.clients.get(clientId);
      if (client?.ws?.readyState === 1) {
        client.ws.send(message);
      }
    }
  }

  async publish(room: string, data: object) {
    const payload = JSON.stringify(data);
    try {
      await redis.publish(room, payload);
    } catch (err: unknown) {
      const message = err instanceof Error ? err.message : "Unknown error";
      logger.error("WebSocket publish failed, falling back to local", {
        room,
        error: message,
      });
      // Fallback to local broadcast if Redis is unavailable
      this.broadcastToRoom(room, payload);
    }
  }
}

export const wsManager = new WebSocketManager();

import { useEffect, useRef, useState } from "react";
import { useAuthStore } from "@/stores/auth";

export function useWebSocket() {
  const wsRef = useRef<WebSocket | null>(null);
  const [isConnected, setIsConnected] = useState(false);
  const { accessToken } = useAuthStore();

  useEffect(() => {
    if (!accessToken) return;

    const ws = new WebSocket("ws://localhost:3001/ws");

    ws.onopen = () => {
      console.log("WebSocket connected");
      setIsConnected(true);

      // Authenticate immediately
      ws.send(JSON.stringify({ type: "auth", token: accessToken }));
    };

    ws.onmessage = (event) => {
      const data = JSON.parse(event.data);
      console.log("WebSocket message:", data);

      // Dispatch custom events for components to listen to
      window.dispatchEvent(
        new CustomEvent("ws-message", { detail: data })
      );
    };

    ws.onerror = (error) => {
      console.error("WebSocket error:", error);
    };

    ws.onclose = () => {
      console.log("WebSocket disconnected");
      setIsConnected(false);
    };

    wsRef.current = ws;

    return () => {
      ws.close();
    };
  }, [accessToken]);

  return { isConnected, ws: wsRef.current };
}

import { wsManager } from "../ws/manager.js";
import { db } from "@restai/db";
import { orders } from "@restai/db/schema";

app.post("/orders", async (c) => {
  const data = await c.req.json();
  
  // Create order in database
  const [order] = await db.insert(orders).values(data).returning();

  // Broadcast to branch room
  await wsManager.publish(`branch:${order.branch_id}`, {
    type: "order:created",
    orderId: order.id,
    orderNumber: order.order_number,
    status: order.status,
    total: order.total,
    timestamp: Date.now(),
  });

  return c.json(order, 201);
});

app.patch("/orders/:id/status", async (c) => {
  const orderId = c.req.param("id");
  const { status } = await c.req.json();

  // Update order
  const [updated] = await db
    .update(orders)
    .set({ status, updated_at: new Date() })
    .where(eq(orders.id, orderId))
    .returning();

  // Broadcast status change
  await wsManager.publish(`branch:${updated.branch_id}`, {
    type: "order:updated",
    orderId: updated.id,
    status: updated.status,
    timestamp: Date.now(),
  });

  // If table session exists, notify table
  if (updated.table_session_id) {
    const session = await db.query.tableSessions.findFirst({
      where: eq(tableSessions.id, updated.table_session_id),
    });

    if (session) {
      await wsManager.publish(`table:${session.table_id}`, {
        type: "order:updated",
        orderId: updated.id,
        status: updated.status,
        timestamp: Date.now(),
      });
    }
  }

  return c.json(updated);
});