File size: 4,196 Bytes
00a912e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
import type { ClientToServerEvents, ServerToClientEvents } from '@music-together/shared'
import { SERVER_URL } from './config'

interface ConnectionEvents {
  connect: () => void
  disconnect: () => void
  connect_error: (error: Error) => void
}

type SocketInboundEvents = ServerToClientEvents & ConnectionEvents
type StoredHandler = (...args: unknown[]) => void

export type TypedSocket = {
  connected: boolean
  connect: () => void
  disconnect: () => void
  on: <E extends keyof SocketInboundEvents>(event: E, handler: SocketInboundEvents[E]) => void
  off: <E extends keyof SocketInboundEvents>(event: E, handler: SocketInboundEvents[E]) => void
  emit: <E extends keyof ClientToServerEvents>(event: E, ...args: Parameters<ClientToServerEvents[E]>) => void
}

let socket: TypedSocket | null = null

export function getSocket(): TypedSocket {
  if (!socket) {
    socket = createWebSocket()
  }
  return socket
}

export function connectSocket(): TypedSocket {
  const s = getSocket()
  s.connect()
  return s
}

export function disconnectSocket(): void {
  if (socket) {
    socket.disconnect()
    socket = null
  }
}

export function waitForConnect(): Promise<TypedSocket> {
  const s = getSocket()
  if (s.connected) return Promise.resolve(s)
  return new Promise((resolve) => {
    const handler = () => {
      s.off('connect', handler)
      resolve(s)
    }
    s.on('connect', handler)
    s.connect()
  })
}

function createWebSocket(): TypedSocket {
  let ws: WebSocket | null = null
  let connected = false
  const handlers = new Map<string, Set<StoredHandler>>()
  let reconnectTimer: ReturnType<typeof setTimeout> | null = null
  let shouldReconnect = true

  const connect = () => {
    shouldReconnect = true
    if (ws && (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING)) {
      return
    }

    const wsUrl = SERVER_URL.replace(/^http/, 'ws') + '/ws'
    ws = new WebSocket(wsUrl)

    ws.onopen = () => {
      connected = true
      const connectHandlers = handlers.get('connect')
      if (connectHandlers) {
        for (const h of connectHandlers) h()
      }
    }

    ws.onmessage = (event) => {
      try {
        const msg = JSON.parse(event.data)
        if (msg && typeof msg.event === 'string') {
          const eventHandlers = handlers.get(msg.event)
          if (eventHandlers) {
            for (const h of eventHandlers) h(msg.data)
          }
        }
      } catch (err) {
        console.error('Failed to parse WebSocket message', err)
      }
    }

    ws.onclose = () => {
      connected = false
      const disconnectHandlers = handlers.get('disconnect')
      if (disconnectHandlers) {
        for (const h of disconnectHandlers) h()
      }

      if (shouldReconnect) {
        reconnectTimer = setTimeout(() => {
          connect()
        }, 2000)
      }
    }

    ws.onerror = (err) => {
      console.error('WebSocket error', err)
      const errorHandlers = handlers.get('connect_error')
      if (errorHandlers) {
        for (const h of errorHandlers) h(new Error('WebSocket error'))
      }
    }
  }

  const disconnect = () => {
    shouldReconnect = false
    if (reconnectTimer) {
      clearTimeout(reconnectTimer)
      reconnectTimer = null
    }
    if (ws) {
      ws.close()
      ws = null
    }
    connected = false
  }

  const on = <E extends keyof SocketInboundEvents>(event: E, handler: SocketInboundEvents[E]) => {
    let set = handlers.get(event as string)
    if (!set) {
      set = new Set()
      handlers.set(event as string, set)
    }
    set.add(handler as unknown as StoredHandler)
  }

  const off = <E extends keyof SocketInboundEvents>(event: E, handler: SocketInboundEvents[E]) => {
    const set = handlers.get(event as string)
    if (set) {
      set.delete(handler as unknown as StoredHandler)
    }
  }

  const emit = <E extends keyof ClientToServerEvents>(event: E, ...args: Parameters<ClientToServerEvents[E]>) => {
    if (ws && ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify({ event, data: args[0] }))
    }
  }

  return {
    get connected() {
      return connected
    },
    connect,
    disconnect,
    on,
    off,
    emit,
  }
}