embody/src/websocket-server.ts

A digital consciousness training environment

import * as http from 'http'
import WebSocket from 'ws'

import * as uuid from 'uuid'

import * as env from './env'

import {
  RobotUsecase,
  RobotUsecaseOptions
} from './usecase'

import {
  RobotManager,
  RobotManagerOptions
} from './manager'

import {
  WebSocketApi,
  WebSocketApiOptions
} from './interface/websocket-api'

import {
  RobotJSProvider,
  RobotJSProviderOptions
} from './provider'

function init (envConfig: any, server?: http.Server): WebSocket.Server {

  // Initialize robotjs as the robot provider.
  const robotjsProviderOptions = new RobotJSProviderOptions()
  const robotjsProvider = new RobotJSProvider(robotjsProviderOptions)

  // Initialize the robot manager.
  const robotManagerOptions = new RobotManagerOptions()
  const robotManager = new RobotManager(robotManagerOptions)
  robotManager.desktopRobot = robotjsProvider

  // Initialize the robot usecase.
  const robotUsecaseOptions = new RobotUsecaseOptions()
  const robotUsecase = new RobotUsecase(robotUsecaseOptions)
  robotUsecase.robotManager = robotManager

  const webSocketApiOptions = new WebSocketApiOptions()

  const w = new WebSocketApi(webSocketApiOptions)
  w.setRobotUsecase(robotUsecase)

  const wss = new WebSocket.Server({ server })

  wss.on('connection', (ws) => {

    const clientID: string = uuid()

    console.log(`WebSocket connection open!`, clientID)

    ws.send(JSON.stringify({
      message: 'WebSocket connection open!',
      clientID
    }))

    ws.on('message', (message) => {
      let data: any

      try {
        data = JSON.parse(message.valueOf() as any)
      } catch (err) {
        console.log(`Error receiving message: ${err}`)
        data = {}
      }

      console.log('Received message: ', data)

      switch (data['type']) {
        case 'start-record-screen': {
          w.robotApi.startRecordingScreen(clientID, sendDataCallback)
          break
        }
        case 'stop-record-screen': {
          w.robotApi.stopRecordingScreen(clientID)
          break
        }
        default:
          console.log(`Unknown message type: ${data['type']}`)
          ws.send(`Unknown message type: ${data['type']}`)
      }
    })

    ws.on('close', () => {
      console.log(`WebSocket connection closed.`, clientID)
      w.robotApi.stopRecordingScreen(clientID)
    })

    ws.on('error', (err) => {
      console.log(`WebSocket error: ${err}`)
    })

    const sendDataCallback = (data: any) => {
      // console.log('Sending message...')
      try {
        ws.send(data)
      } catch (err) {
        console.log(`Error sending data: ${err}`)
      }
    }
  })

  return wss
}

export default {
  init
}