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

import Control.Concurrent
import Control.Concurrent.STM
import Control.Monad

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
    isReliableTransport _ = True

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


startDefaultWebSocketConnections :: Server -> IO ()
startDefaultWebSocketConnections server = do
    peerVar <- newEmptyMVar
    let doStart = do
            startClient server "a.discovery.erebosprotocol.net" 443 "" $ \conn -> do
                putMVar peerVar =<< serverPeerCustom server conn
                void $ forkIO $ forever $ do
                    msg <- receiveMessage conn
                    receivedFromCustomAddress server conn msg

    pchan <- getNextPeerChangeChan server
    void $ forkIO $ void $ forever $ do
        peer <- atomically $ readTChan pchan
        isPeerDropped peer >>= \case
            True -> do
                wsPeer <- takeMVar peerVar
                if peer == wsPeer
                  then doStart
                  else putMVar peerVar wsPeer
            False -> return ()

    doStart


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 ()