stremio/apps/master-api/src/websocket/manager.ts

import type { WebSocket } from "ws";
import type { ControlMessage, ControlMessageType } from "@media-cluster/protocol";
import { createMessage } from "@media-cluster/protocol";
import type { CreatePlaybackSessionPayload } from "@media-cluster/shared-types";
import { prisma } from "../database/client";
import { authenticateNode } from "../nodes/routes";
import { LibrarySyncService } from "../media/sync";
import { MetadataService } from "../metadata/service";
import type { Config } from "../config";
import type {
  FullLibrarySyncPayload,
  HeartbeatPayload,
  HelloPayload,
  LibraryEventPayload,
  PlaybackSessionEndedPayload,
  CreatePlaybackSessionAckPayload,
} from "@media-cluster/protocol";

interface NodeConnection {
  nodeId: string;
  socket: WebSocket;
  lastHeartbeat: Date;
}

interface PendingAck {
  resolve: (ok: boolean) => void;
  timer: ReturnType<typeof setTimeout>;
}

interface SyncState {
  syncId: string;
  fileIds: Set<string>;
}

class NodeConnectionManager {
  private connections = new Map<string, NodeConnection>();
  private librarySync: LibrarySyncService | null = null;
  private config: Config | null = null;
  private offlineCheckInterval: ReturnType<typeof setInterval> | null = null;
  private pendingSessionAcks = new Map<string, PendingAck>();
  private syncState = new Map<string, SyncState>();
  private syncQueues = new Map<string, Promise<void>>();

  init(config: Config): void {
    this.config = config;
    const metadata = new MetadataService(config.TMDB_API_KEY);
    this.librarySync = new LibrarySyncService(metadata);

    this.offlineCheckInterval = setInterval(() => {
      void this.checkOfflineNodes();
    }, 30_000);
  }

  async handleConnection(socket: WebSocket): Promise<void> {
    let nodeId: string | null = null;
    let authenticated = false;

    const authTimeout = setTimeout(() => {
      if (!authenticated) socket.close(4001, "Authentication timeout");
    }, 30_000);

    socket.on("message", async (data) => {
      try {
        const message = JSON.parse(data.toString()) as ControlMessage;
        await this.handleMessage(socket, message, (id) => {
          nodeId = id;
          authenticated = true;
          clearTimeout(authTimeout);
        });
      } catch (err) {
        console.error("WebSocket message error:", err);
        this.send(socket, createMessage("ERROR", { message: "Invalid message" }));
      }
    });

    socket.on("close", () => {
      if (nodeId) {
        this.connections.delete(nodeId);
        this.syncState.delete(nodeId);
        void prisma.node.update({
          where: { id: nodeId },
          data: { status: "OFFLINE" },
        });
        void prisma.mediaFile.updateMany({
          where: { nodeId },
          data: { available: false },
        });
        console.log(`Node disconnected: ${nodeId}`);
      }
    });
  }

  private async handleMessage(
    socket: WebSocket,
    message: ControlMessage,
    onAuth: (nodeId: string) => void
  ): Promise<void> {
    if (message.protocolVersion !== 1) {
      this.send(
        socket,
        createMessage("ERROR", {
          code: "UNSUPPORTED_PROTOCOL",
          message: `Unsupported protocol version: ${message.protocolVersion}`,
        })
      );
      socket.close(4002, "Unsupported protocol");
      return;
    }

    switch (message.type) {
      case "HELLO":
        await this.handleHello(socket, message.payload as HelloPayload, onAuth);
        break;
      case "HEARTBEAT":
        await this.handleHeartbeat(message.payload as HeartbeatPayload);
        this.send(socket, createMessage("HEARTBEAT_ACK", {}));
        break;
      case "LIBRARY_EVENT":
        await this.handleLibraryEvent(
          message.payload as LibraryEventPayload,
          socket
        );
        break;
      case "FULL_LIBRARY_SYNC":
        await this.handleFullSync(message.payload as FullLibrarySyncPayload, socket);
        break;
      case "PLAYBACK_SESSION_ENDED":
        await this.handlePlaybackEnded(message.payload as PlaybackSessionEndedPayload);
        break;
      case "CREATE_PLAYBACK_SESSION_ACK":
        this.handlePlaybackSessionAck(message.payload as CreatePlaybackSessionAckPayload);
        break;
      default:
        this.send(socket, createMessage("ERROR", { message: `Unknown type: ${message.type}` }));
    }
  }

  private async handleHello(
    socket: WebSocket,
    payload: HelloPayload,
    onAuth: (nodeId: string) => void
  ): Promise<void> {
    const valid = await authenticateNode(payload.nodeId, payload.apiKey);
    if (!valid) {
      this.send(
        socket,
        createMessage("HELLO_ACK", {
          accepted: false,
          reason: "Invalid credentials",
          heartbeatIntervalSeconds: 30,
          fullScanIntervalSeconds: 0,
          fullScanInterval: "off",
          fullScanAt: "03:30",
        })
      );
      socket.close(4003, "Authentication failed");
      return;
    }

    this.syncState.delete(payload.nodeId);
    this.connections.set(payload.nodeId, {
      nodeId: payload.nodeId,
      socket,
      lastHeartbeat: new Date(),
    });
    onAuth(payload.nodeId);

    await prisma.node.update({
      where: { id: payload.nodeId },
      data: {
        status: "ONLINE",
        lastSeen: new Date(),
        version: payload.softwareVersion,
        architecture: payload.architecture,
        hostname: payload.hostname,
        // Do not overwrite publicStreamUrl on HELLO — admin may correct typos
        // (e.g. node01 vs node1) without recreating the Unraid container env.
      },
    });

    // Restore availability for previously indexed files until full sync arrives
    await prisma.mediaFile.updateMany({
      where: { nodeId: payload.nodeId },
      data: { available: true },
    });

    console.log(`Node connected: ${payload.nodeId} (${payload.hostname})`);

    let node = await prisma.node.findUnique({
      where: { id: payload.nodeId },
      include: { scanRoots: { include: { shelf: true } } },
    });

    const { ensureDefaultShelves, DEFAULT_MOVIE_SHELF_ID, DEFAULT_SERIES_SHELF_ID, getNodeScanRootsPayload, deriveLegacyPaths } =
      await import("../library/shelves");
    await ensureDefaultShelves();

    // Seed master path lists from node env/config when admin has not set them yet.
    if (
      node &&
      (node.scanRoots?.length ?? 0) === 0 &&
      (node.moviesPaths?.length ?? 0) === 0 &&
      (node.seriesPaths?.length ?? 0) === 0 &&
      ((payload.moviesPaths?.length ?? 0) > 0 || (payload.seriesPaths?.length ?? 0) > 0)
    ) {
      const rows: Array<{ nodeId: string; path: string; shelfId: string }> = [];
      for (const p of payload.moviesPaths ?? []) {
        rows.push({ nodeId: payload.nodeId, path: p, shelfId: DEFAULT_MOVIE_SHELF_ID });
      }
      for (const p of payload.seriesPaths ?? []) {
        rows.push({ nodeId: payload.nodeId, path: p, shelfId: DEFAULT_SERIES_SHELF_ID });
      }
      if (rows.length) {
        await prisma.nodeScanRoot.createMany({ data: rows });
        const legacy = deriveLegacyPaths(
          rows.map((r) => ({
            path: r.path,
            kind: r.shelfId === DEFAULT_SERIES_SHELF_ID ? ("series" as const) : ("movie" as const),
          }))
        );
        node = await prisma.node.update({
          where: { id: payload.nodeId },
          data: { moviesPaths: legacy.moviesPaths, seriesPaths: legacy.seriesPaths },
          include: { scanRoots: { include: { shelf: true } } },
        });
      }
    } else if (node && (node.scanRoots?.length ?? 0) === 0 && ((node.moviesPaths?.length ?? 0) > 0 || (node.seriesPaths?.length ?? 0) > 0)) {
      // Migrate legacy arrays → scan roots once
      const rows: Array<{ nodeId: string; path: string; shelfId: string }> = [];
      for (const p of node.moviesPaths) {
        rows.push({ nodeId: node.id, path: p, shelfId: DEFAULT_MOVIE_SHELF_ID });
      }
      for (const p of node.seriesPaths) {
        rows.push({ nodeId: node.id, path: p, shelfId: DEFAULT_SERIES_SHELF_ID });
      }
      if (rows.length) {
        await prisma.nodeScanRoot.createMany({ data: rows });
        node = await prisma.node.findUnique({
          where: { id: payload.nodeId },
          include: { scanRoots: { include: { shelf: true } } },
        });
      }
    }

    const fullScanInterval = node?.fullScanInterval || "off";
    const fullScanAt = node?.fullScanAt || "03:30";
    const legacySeconds =
      fullScanInterval === "off"
        ? 0
        : (() => {
            const m = fullScanInterval.match(/^(\d+)h$/i);
            return m ? parseInt(m[1], 10) * 3600 : 0;
          })();

    const rootsPayload = await getNodeScanRootsPayload(payload.nodeId);
    const legacy = deriveLegacyPaths(rootsPayload);

    this.send(
      socket,
      createMessage("HELLO_ACK", {
        accepted: true,
        heartbeatIntervalSeconds: this.config?.NODE_HEARTBEAT_INTERVAL_SECONDS ?? 30,
        fullScanIntervalSeconds: legacySeconds,
        fullScanInterval,
        fullScanAt,
        moviesPaths: legacy.moviesPaths,
        seriesPaths: legacy.seriesPaths,
        scanRoots: rootsPayload.map((r) => ({
          path: r.path,
          shelfId: r.shelfId,
          kind: r.kind,
        })),
      })
    );
  }

  private async handleHeartbeat(payload: HeartbeatPayload): Promise<void> {
    const conn = this.connections.get(payload.nodeId);
    if (conn) conn.lastHeartbeat = new Date();

    await prisma.node.update({
      where: { id: payload.nodeId },
      data: {
        status: "ONLINE",
        lastSeen: new Date(),
        totalStorage: BigInt(payload.storageTotalBytes),
        freeStorage: BigInt(payload.storageFreeBytes),
        activeStreams: payload.activeStreams,
        currentBandwidth: BigInt(payload.currentBytesPerSec),
        libraryFileCount: payload.libraryFileCount,
        version: payload.softwareVersion,
      },
    });
  }

  private async handleLibraryEvent(
    payload: LibraryEventPayload,
    socket: WebSocket
  ): Promise<void> {
    const conn = [...this.connections.values()].find((c) => c.socket === socket);
    if (!conn || !this.librarySync) return;

    // Don't request another full sync while one is already in progress
    if (this.syncState.has(conn.nodeId)) {
      await this.librarySync.processEvents(conn.nodeId, payload.events);
      this.send(socket, createMessage("LIBRARY_EVENT_ACK", { accepted: true }));
      return;
    }

    const node = await prisma.node.findUnique({ where: { id: conn.nodeId } });
    if (!node) return;

    const expectedRevision = node.lastRevision + 1;
    const minRevision = Math.min(...payload.events.map((e) => e.nodeRevision));

    if (payload.events.length > 0 && minRevision > expectedRevision) {
      this.send(socket, createMessage("FULL_LIBRARY_SYNC", { full: true }));
      return;
    }

    await this.librarySync.processEvents(conn.nodeId, payload.events);
    this.send(socket, createMessage("LIBRARY_EVENT_ACK", { accepted: true }));
  }

  private async handleFullSync(
    payload: FullLibrarySyncPayload,
    socket: WebSocket
  ): Promise<void> {
    const conn = [...this.connections.values()].find((c) => c.socket === socket);
    if (!conn || !this.librarySync) return;

    const syncId = payload.syncId ?? `legacy-${Date.now()}`;
    const isLast = payload.isLast ?? true;
    const batchIndex = payload.batchIndex ?? 0;

    let state = this.syncState.get(conn.nodeId);
    if (!state || state.syncId !== syncId) {
      state = { syncId, fileIds: new Set() };
      this.syncState.set(conn.nodeId, state);
    }

    // ACK immediately so the node WS stays free for playback / next batches.
    // Metadata matching can take seconds per batch and must not block the socket.
    this.send(
      socket,
      createMessage("FULL_LIBRARY_SYNC_ACK", {
        accepted: true,
        syncId: payload.syncId,
        batchIndex: payload.batchIndex,
      })
    );

    const prev = this.syncQueues.get(conn.nodeId) ?? Promise.resolve();
    const job = prev
      .then(async () => {
        if (!this.librarySync) return;
        await this.librarySync.processFullSync(
          conn.nodeId,
          payload.nodeRevision,
          payload.files,
          { isLast, fileIdsSoFar: state!.fileIds }
        );

        if (isLast) {
          console.log(
            `Full library sync complete for ${conn.nodeId}: ${state!.fileIds.size} files` +
              (payload.batchCount != null ? ` (${payload.batchCount} batches)` : "")
          );
          this.syncState.delete(conn.nodeId);
        } else if (batchIndex % 10 === 0) {
          console.log(
            `Full library sync progress for ${conn.nodeId}: batch ${(batchIndex ?? 0) + 1}/${payload.batchCount ?? "?"}`
          );
        }
      })
      .catch((err) => {
        console.error(`Full sync failed for ${conn.nodeId}:`, err);
        this.syncState.delete(conn.nodeId);
      });
    this.syncQueues.set(conn.nodeId, job);
  }

  private handlePlaybackSessionAck(payload: CreatePlaybackSessionAckPayload): void {
    const pending = this.pendingSessionAcks.get(payload.sessionId);
    if (!pending) return;
    clearTimeout(pending.timer);
    this.pendingSessionAcks.delete(payload.sessionId);
    pending.resolve(payload.ok);
  }

  private async handlePlaybackEnded(payload: PlaybackSessionEndedPayload): Promise<void> {
    await prisma.playbackSession.updateMany({
      where: { id: payload.sessionId },
      data: { status: "COMPLETED" },
    });
  }

  sendToNode(nodeId: string, type: ControlMessageType, payload: unknown): boolean {
    const conn = this.connections.get(nodeId);
    if (!conn || conn.socket.readyState !== 1) return false;
    this.send(conn.socket, createMessage(type, payload));
    return true;
  }

  sendRescan(nodeId: string, full: boolean): boolean {
    return this.sendToNode(nodeId, "RESCAN", { full });
  }

  sendConfigUpdate(
    nodeId: string,
    payload: {
      fullScanInterval?: string;
      fullScanAt?: string;
      moviesPaths?: string[];
      seriesPaths?: string[];
      scanRoots?: Array<{ path: string; shelfId: string; kind: "movie" | "series" }>;
    }
  ): boolean {
    return this.sendToNode(nodeId, "CONFIG_UPDATE", payload);
  }

  sendRestart(nodeId: string, reason = "admin"): boolean {
    return this.sendToNode(nodeId, "RESTART", { reason });
  }

  createPlaybackSession(nodeId: string, payload: CreatePlaybackSessionPayload): boolean {
    return this.sendToNode(nodeId, "CREATE_PLAYBACK_SESSION", payload);
  }

  /** Push session to node and wait for ACK (or timeout). */
  async createPlaybackSessionAcked(
    nodeId: string,
    payload: CreatePlaybackSessionPayload,
    timeoutMs = 5000
  ): Promise<boolean> {
    if (!this.isOnline(nodeId)) return false;

    const acked = new Promise<boolean>((resolve) => {
      const timer = setTimeout(() => {
        // Fail closed: never hand Stremio a node URL without a confirmed session.
        this.pendingSessionAcks.delete(payload.sessionId);
        resolve(false);
      }, timeoutMs);
      this.pendingSessionAcks.set(payload.sessionId, { resolve, timer });
    });

    const pushed = this.createPlaybackSession(nodeId, payload);
    if (!pushed) {
      const pending = this.pendingSessionAcks.get(payload.sessionId);
      if (pending) {
        clearTimeout(pending.timer);
        this.pendingSessionAcks.delete(payload.sessionId);
      }
      return false;
    }

    return acked;
  }

  revokePlaybackSession(nodeId: string, sessionId: string): boolean {
    return this.sendToNode(nodeId, "REVOKE_PLAYBACK_SESSION", { sessionId });
  }

  disconnect(nodeId: string): void {
    const conn = this.connections.get(nodeId);
    if (conn) {
      conn.socket.close(4004, "Revoked");
      this.connections.delete(nodeId);
      this.syncState.delete(nodeId);
    }
  }

  isOnline(nodeId: string): boolean {
    const conn = this.connections.get(nodeId);
    return !!conn && conn.socket.readyState === 1;
  }

  private send(socket: WebSocket, message: ControlMessage): void {
    if (socket.readyState === 1) {
      socket.send(JSON.stringify(message));
    }
  }

  private async checkOfflineNodes(): Promise<void> {
    if (!this.config) return;
    const threshold = new Date(
      Date.now() - this.config.NODE_OFFLINE_THRESHOLD_SECONDS * 1000
    );

    const stale = await prisma.node.findMany({
      where: {
        status: "ONLINE",
        lastSeen: { lt: threshold },
        revoked: false,
      },
      select: { id: true },
    });

    for (const node of stale) {
      if (this.connections.has(node.id)) {
        this.disconnect(node.id);
      }
      await prisma.node.update({
        where: { id: node.id },
        data: { status: "OFFLINE" },
      });
      await prisma.mediaFile.updateMany({
        where: { nodeId: node.id },
        data: { available: false },
      });
    }
  }

  shutdown(): void {
    if (this.offlineCheckInterval) clearInterval(this.offlineCheckInterval);
    for (const conn of this.connections.values()) {
      conn.socket.close(1001, "Server shutting down");
    }
    this.connections.clear();
    this.syncState.clear();
  }
}

export const nodeConnectionManager = new NodeConnectionManager();