import { Server } from 'ws'
import { findKeyAndIV } from '@ideadesignmedia/develop/pwa/src/auth'
import { decrypt, encrypt, randomString } from '@ideadesignmedia/develop/pwa/src/encryption'

const wsPort = parseInt(process.env.WS_PORT || '0') || 2800
const wss = new Server({ port: wsPort }) // Create WebSocket server
wss.on('listening', () => console.log(`WebSocket server listening on port ${wsPort}`))
const socketKeys = new Map() // Map of WebSocket connections to their encryption keys
const authTimeout = new Map() // Map of WebSocket connections to their authentication timeouts
export const socketRecording = new Map() // Map of WebSocket connections to their recordings
const patternResponse = (type: string = 'response', data: any = {}, error?: string, responseCode?: any) => ({type, data, error, responseCode })
const encryptedResponse = (e: string, { key, iv }: { key: string, iv: string } = { key: '', iv: '' }) => {
  if (!key || !iv) return e // Don't encrypt if there's no key or IV
  return `~e~${encrypt(e, key, iv)}` // Encrypt a response before sending it to the client
}
wss.on('connection', function connection(ws, req) {
  authTimeout.set(ws, setTimeout(() => { // Set a timeout for authentication
    ws.reply(patternResponse('auth', {}, 'Timeout'))
    ws.close()
  }, 5000))
  ws.reply = (e) => { // Send a reply to the client
    ws.send(encryptedResponse(typeof e === 'string' ? e : JSON.stringify(e), socketKeys.get(ws)))
  }
  ws.on('close', function close() {
    if (authTimeout.has(ws)) {
      clearTimeout(authTimeout.get(ws))
      authTimeout.delete(ws)
    }
    if (socketKeys.has(ws)) {
      const { token } = socketKeys.get(ws)
      if (socketRecording.has(token)) {
        socketRecording.get(token).stream.end()
        socketRecording.delete(token)
      }
      socketKeys.delete(ws)
    }
  })
  ws.on('error', e => console.error(`WS ERROR: ${JSON.stringify(e)}`))
  ws.on('message', async function incoming(message) {
    let that = message instanceof Buffer ? message.toString() : message
    if (typeof that === 'string') {
      if (that.startsWith('~e~')) {
        try {
          that = decrypt(that.substring(3), socketKeys.get(ws).key, socketKeys.get(ws).iv)
        } catch {
          ws.reply(patternResponse('auth', {}, 'Unauthorized'))
          ws.close()
          return
        }
      }
      try {
        that = JSON.parse(that)
      } catch {
        that = null
      }
      if (!that) {
        that = message instanceof Buffer ? message.toString() : message
      }
    }
    if (typeof that === 'object' || !(that instanceof Buffer)) {
      const { type, data } = that
      if (!socketKeys.has(ws) && type !== 'auth') {
        ws.reply(patternResponse('auth', {}, 'Unauthorized'))
        ws.close()
        return
      }
      switch (type) {
        case 'ping': {
          ws.reply({ type: 'pong', data })
          break
        }
        case 'auth': {
          const { token } = data
          findKeyAndIV(token).then(({ key, iv }) => {
            clearTimeout(authTimeout.get(ws))
            authTimeout.delete(ws)
            socketKeys.set(ws, { key, iv, token })
            ws.reply(patternResponse('auth', {success: true}))
          }).catch((e) => {
            console.error(e)
            ws.reply(patternResponse('auth', {}, 'Unauthorized'))
            ws.close()
          })
        }
      }
    } else {
      ws.reply(patternResponse('auth', {}, 'Unauthorized'))
      ws.close()
    }
  })
})
export default wss;