import { useCallback, useEffect, useRef, useState } from 'react'
import { useAppContext } from '../app';
import AdvancedWebSocket from '../utils/advanced-websocket';
import { decrypt } from '../utils/encryption';
export function useWS(url: string, onMessage: (data: any, ws: AdvancedWebSocket | null) => void) {
    const { state } = useAppContext() as any
    const tokenRef = useRef<{ token: string, key: string, iv: string } | null>(null)
    const [websocket, setWS] = useState<AdvancedWebSocket | null>(null)
    const [refreshSocket, setRefresh] = useState(0)
    const open = useCallback(() => {
        if (websocket) websocket.close(4900)
        setRefresh(r => r + 1)
    }, [websocket])
    const close = useCallback(() => {
        if (websocket) websocket.close(4900)
        setWS(null)
    }, [websocket])
    useEffect(() => {
        if (state.token) {
            tokenRef.current = { token: state.token, key: state.key, iv: state.iv }
        } else if (!state.token) {
            tokenRef.current = null
        }
    }, [state.token, state.key, state.iv])
    const websocketRef = useRef<AdvancedWebSocket | null>(null)
    useEffect(() => {
        return () => {
            if (websocketRef.current && websocketRef.current.readyState === WebSocket.OPEN) {
                websocketRef.current.close(4900)
            }
        }
    }, [])
    useEffect(() => {
        if (websocketRef.current) {
            if (!url) setWS(ws => {
                if (ws) ws.close(4900)
                websocketRef.current = null
                return null
            })
        }
        if (!websocketRef.current && url) {
            setRefresh(r => r + 1)
        }
    }, [url])
    useEffect(() => {
        if (!tokenRef.current) return setWS(ws => {
            if (ws) {
                ws.close(4900)
                websocketRef.current = null
            }
            return null
        })
        if (!websocketRef.current && url) {
            const createSocket = () => {
                let ws: AdvancedWebSocket | null
                try {
                    ws = new AdvancedWebSocket(url, tokenRef.current?.key, tokenRef.current?.iv)
                } catch (e) {
                    console.error(e)
                    return null
                }
                if (!ws) return null
                let lastPing = Date.now()
                const heartbeat = setInterval(() => {
                    lastPing = Date.now()
                    ws?.sendEncrypted(JSON.stringify({ type: 'ping' }))
                }, 30000)
                ws.onopen = () => {
                    if (!tokenRef.current) {
                        ws?.close(4900)
                        return
                    }
                    ws?.sendData(JSON.stringify({ type: 'auth', data: { token: tokenRef.current.token } }))
                }
                ws.onmessage = (data: any) => {
                    let o: string | object | null = null
                    if (tokenRef.current && (typeof data === 'string' || typeof data.data === 'string') && ((data?.data || data).startsWith('~e~'))) {
                        try {
                            o = decrypt((data?.data || data).substring(3), tokenRef.current.key, tokenRef.current.iv)
                        } catch (e) { o = null, console.error(e) }
                    }
                    if (!o) {
                        try {
                            o = JSON.parse(data)
                        } catch {
                            try {
                                o = JSON.parse(data.data)
                            } catch { o = null }
                        }
                        if (o) {
                            if ((o as { type: string }).type === 'pong') {
                                console.log(`WebSocket Latency: ${Date.now() - lastPing}ms`)
                                return
                            }
                            onMessage(o, ws)
                        } else if (data.data) {
                            onMessage(data.data, ws)
                        }
                    } else {
                        if (typeof o === 'string' && (o.startsWith('{') || o.startsWith('['))) {
                            try {
                                o = JSON.parse(o)
                            } catch {
                                onMessage(o, ws)
                                o = null
                            }
                            if (o) {
                                if ((o as { type: string }).type === 'pong') {
                                    console.log(`WebSocket Latency: ${Date.now() - lastPing}ms`)
                                    return
                                }
                                onMessage(o, ws)
                            }
                        } else {
                            onMessage(o, ws)
                        }
                    }
                }
                const closeSocket = (event: CloseEvent) => {
                    clearInterval(heartbeat)
                    if (tokenRef.current && event && event?.code !== 4900) { // 4900 is the code for a manual close
                        setTimeout(() => {
                            websocketRef.current = createSocket()
                            setWS(websocketRef.current)
                        }, 1000)
                        return false
                    } else {
                        websocketRef.current = null
                        setWS(null)
                        return true
                    }
                }
                ws.onclose = closeSocket as any
                return ws
            }
            websocketRef.current = createSocket()
            setWS(websocketRef.current)
        }
    }, [onMessage, refreshSocket])
    return { ws: websocket, close, open }
}
export default useWS;