unit WebSocket.Manager;

interface

uses
  System.Classes,
  System.SysUtils,
  System.JSON,
  System.Generics.Collections,
  VCL.TMSFNCWebSocketServer,
  VCL.TMSFNCWebSocketCommon;

const
  WEBSOCKET_PORT = 8091;

type
  TConnectedClientSnapshot = record
    ConnectionId: string;
    UserId: string;
    ConnectedAt: TDateTime;
  end;

  TConnectedClient = class
  private
    FConnectionId: string;
    FUserId: string;
    FConnectedAt: TDateTime;
    FConnection: TTMSFNCWebSocketServerConnection;
  public
    property ConnectionId: string read FConnectionId write FConnectionId;
    property UserId: string read FUserId write FUserId;
    property ConnectedAt: TDateTime read FConnectedAt write FConnectedAt;
    property Connection: TTMSFNCWebSocketServerConnection
      read FConnection write FConnection;
  end;

  TClientsChangedEvent = procedure of object;

  TWebSocketManager = class
  private
    FServer: TTMSFNCWebSocketServer;
    FClients: TObjectList<TConnectedClient>;
    FClientsLock: TObject;
    FOnClientsChanged: TClientsChangedEvent;

    procedure NotifyClientsChanged;
    procedure HandshakeResponseSent(Sender: TObject;
      AConnection: TTMSFNCWebSocketServerConnection);
    procedure MessageReceived(Sender: TObject;
      AConnection: TTMSFNCWebSocketConnection; const AMessage: string);
    procedure ClientDisconnected(Sender: TObject;
      AConnection: TTMSFNCWebSocketConnection);
  public
    constructor Create;
    destructor Destroy; override;

    procedure Start;
    procedure Stop;

    // Send AMessage to every currently connected and handshaked client.
    // Snapshots the connection list under lock before sending so that the lock
    // is never held while performing I/O.
    procedure Broadcast(const AMessage: string);

    function GetClientSnapshots: TArray<TConnectedClientSnapshot>;

    property OnClientsChanged: TClientsChangedEvent
      read FOnClientsChanged write FOnClientsChanged;
  end;

implementation

uses
  Common.Logging;

constructor TWebSocketManager.Create;
begin
  inherited Create;

  FClientsLock := TObject.Create;
  FClients := TObjectList<TConnectedClient>.Create(True);

  FServer := TTMSFNCWebSocketServer.Create;
  FServer.Port := WEBSOCKET_PORT;
  FServer.UseSSL := False;
  FServer.OnHandshakeResponseSent := HandshakeResponseSent;
  FServer.OnMessageReceived := MessageReceived;
  FServer.OnDisconnect := ClientDisconnected;
end;

destructor TWebSocketManager.Destroy;
begin
  Stop;
  FServer.Free;
  FClients.Free;
  FClientsLock.Free;

  inherited;
end;

procedure TWebSocketManager.Start;
begin
  FServer.Active := True;
end;

procedure TWebSocketManager.Stop;
begin
  FServer.Active := False;
end;

procedure TWebSocketManager.NotifyClientsChanged;
begin
  if Assigned(FOnClientsChanged) then
    FOnClientsChanged;
end;

procedure TWebSocketManager.HandshakeResponseSent(Sender: TObject;
  AConnection: TTMSFNCWebSocketServerConnection);
var
  client: TConnectedClient;
  guid: TGUID;
begin
  CreateGUID(guid);

  client := TConnectedClient.Create;
  client.ConnectionId := GUIDToString(guid);
  client.ConnectedAt := Now;
  client.Connection := AConnection;

  TMonitor.Enter(FClientsLock);
  try
    FClients.Add(client);
  finally
    TMonitor.Exit(FClientsLock);
  end;

  AConnection.OwnsUserData := False;
  AConnection.UserData := client;

  Logger.Log(1, 'WS: client connected [' + client.ConnectionId + ']');
  NotifyClientsChanged;
end;

procedure TWebSocketManager.MessageReceived(Sender: TObject;
  AConnection: TTMSFNCWebSocketConnection; const AMessage: string);
var
  json: TJSONValue;
  messageType, userId, connectionId: string;
  client: TConnectedClient;
begin
  json := TJSONObject.ParseJSONValue(AMessage);
  try
    if not Assigned(json) then
      Exit;

    if not json.TryGetValue<string>('message', messageType) then
      Exit;

    // "identify" lets the client register a human-readable userId for the log.
    if messageType <> 'identify' then
      Exit;

    if not json.TryGetValue<string>('userId', userId) then
      Exit;

    client := TConnectedClient(
      TTMSFNCWebSocketServerConnection(AConnection).UserData);

    if not Assigned(client) then
      Exit;

    TMonitor.Enter(FClientsLock);
    try
      if FClients.IndexOf(client) < 0 then
        Exit;

      client.UserId := userId;
      connectionId := client.ConnectionId;
    finally
      TMonitor.Exit(FClientsLock);
    end;

    Logger.Log(1, 'WS: client identified [' + connectionId + '] userId=' + userId);
    NotifyClientsChanged;
  finally
    json.Free;
  end;
end;

procedure TWebSocketManager.ClientDisconnected(Sender: TObject;
  AConnection: TTMSFNCWebSocketConnection);
var
  serverConn: TTMSFNCWebSocketServerConnection;
  client: TConnectedClient;
  connectionId, userId: string;
begin
  serverConn := TTMSFNCWebSocketServerConnection(AConnection);
  client := TConnectedClient(serverConn.UserData);

  if not Assigned(client) then
    Exit;

  connectionId := client.ConnectionId;
  userId := client.UserId;
  serverConn.UserData := nil;

  TMonitor.Enter(FClientsLock);
  try
    FClients.Remove(client);
  finally
    TMonitor.Exit(FClientsLock);
  end;

  Logger.Log(1, 'WS: client disconnected [' + connectionId + '] userId=' + userId);
  NotifyClientsChanged;
end;

procedure TWebSocketManager.Broadcast(const AMessage: string);
var
  snapshot: TArray<TTMSFNCWebSocketServerConnection>;
  i: Integer;
begin
  // Snapshot connections under lock so we never hold the lock during I/O.
  TMonitor.Enter(FClientsLock);
  try
    SetLength(snapshot, FClients.Count);
    for i := 0 to FClients.Count - 1 do
      snapshot[i] := FClients[i].Connection;
  finally
    TMonitor.Exit(FClientsLock);
  end;

  for i := 0 to Length(snapshot) - 1 do
  try
    snapshot[i].WriteString(AMessage);
  except
    // Ignore disconnected or errored connections; they will be cleaned up
    // via the OnDisconnect event.
  end;
end;

function TWebSocketManager.GetClientSnapshots: TArray<TConnectedClientSnapshot>;
var
  i: Integer;
begin
  TMonitor.Enter(FClientsLock);
  try
    SetLength(Result, FClients.Count);
    for i := 0 to FClients.Count - 1 do
    begin
      Result[i].ConnectionId := FClients[i].ConnectionId;
      Result[i].UserId := FClients[i].UserId;
      Result[i].ConnectedAt := FClients[i].ConnectedAt;
    end;
  finally
    TMonitor.Exit(FClientsLock);
  end;
end;

end.
