module WebSocket (
    Connection,
    startClient,
    sendMessage,
    receiveMessage,
) where

import Control.Concurrent.Chan

import Data.ByteString (ByteString)
import Data.ByteString.Unsafe
import Data.Function
import Data.Unique
import Data.Word

import Erebos.Network

import Foreign.Marshal.Alloc
import Foreign.Ptr

import GHC.Wasm.Prim

import JavaScript
import JavaScript.Val


data Connection = Connection
    { connUnique :: Unique
    , connAddress :: String
    , connJS :: WebSocket
    , connInQueue :: Chan ByteString
    }

instance Eq Connection where
    (==) = (==) `on` connUnique

instance Ord Connection where
    compare = compare `on` connUnique

instance Show Connection where
    show = connAddress

instance PeerAddressType Connection where
    sendBytesToAddress = sendMessage
    connectionToAddressClosed = closeConnection

startClient :: Server -> String -> Int -> String -> (Connection -> IO ()) -> IO ()
startClient server addr port path fun = do
    connUnique <- newUnique
    let connAddress = "wss://" <> addr <> ":" <> show port <> "/" <> path
    connJS <- initWebSocket connAddress
    connInQueue <- newChan
    let conn = Connection {..}

    addEventListener connJS "open" $ \_ -> do
        fun conn

    addEventListener connJS "message" $ \ev -> do
        bytes <- js_get_data $ toJSVal ev
        len <- js_get_byteLength bytes
        ptr <- mallocBytes len
        js_copyBytes ptr bytes
        bs <- unsafePackCStringFinalizer ptr len (free ptr)
        writeChan connInQueue bs

    addEventListener connJS "close" $ \_ -> do
        dropPeerAddress server $ CustomPeerAddress conn

    addEventListener connJS "error" $ \_ -> do
        dropPeerAddress server $ CustomPeerAddress conn


sendMessage :: Connection -> ByteString -> IO ()
sendMessage Connection {..} bs = do
    unsafeUseAsCStringLen bs $ \( ptr, len ) -> do
        js_send (toJSVal connJS) (castPtr ptr) len >>= \case
            0 -> return ()
            1 -> error "websocket is not open"

receiveMessage :: Connection -> IO ByteString
receiveMessage Connection {..} = do
    readChan connInQueue

closeConnection :: Connection -> IO ()
closeConnection Connection {..} = do
    js_close (toJSVal connJS)


newtype WebSocket = WebSocket JSVal
    deriving (FromJSVal, ToJSVal)

instance IsEventTarget WebSocket where
    toEventTarget = fromJSValUnchecked . toJSVal

initWebSocket :: String -> IO WebSocket
initWebSocket addr = WebSocket <$> js_initWebSocket (toJSString addr)

foreign import javascript unsafe "const ws = new WebSocket($1); ws.binaryType = 'arraybuffer'; return ws"
    js_initWebSocket :: JSString -> IO JSVal

foreign import javascript unsafe "if ($1.readyState == WebSocket.OPEN) { $1.send(new Uint8Array(globalThis.wasi_memory.buffer, $2, $3)); return 0; } else { return 1; }"
    js_send :: JSVal -> Ptr Word8 -> Int -> IO Int

foreign import javascript unsafe "$1.close()"
    js_close :: JSVal -> IO ()

foreign import javascript unsafe "$1.data"
    js_get_data :: JSVal -> IO JSVal

foreign import javascript unsafe "$1.byteLength"
    js_get_byteLength :: JSVal -> IO Int

foreign import javascript unsafe "new Uint8Array(globalThis.wasi_memory.buffer, $1, $2.byteLength).set(new Uint8Array($2))"
    js_copyBytes :: Ptr Word8 -> JSVal -> IO ()