{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -fno-warn-unused-binds -fno-warn-unused-matches -fno-warn-missing-signatures #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}

import qualified Network.WebSockets as WS
import qualified Data.Aeson as A
import Data.Text (Text)
import Control.Concurrent (forkFinally)
import GHC.Generics (Generic)
import Data.Aeson (ToJSON (toJSON), FromJSON (parseJSON), (.:))
import Data.ByteString.Lazy (LazyByteString)
import Control.Concurrent.MVar (newEmptyMVar, putMVar, MVar, readMVar)
import GHC.Stack (HasCallStack)


data ClientMessage = ClientMessage
    { targetId :: !(Maybe Text)
    , body :: !ClientMessageBody
    }
    deriving (Eq, Show, Generic)
    deriving anyclass (ToJSON, FromJSON)
data ClientMessageBody
    = ClientSdpMessage !Sdp
    | ClientIceMessage !A.Value
    | JoinRoomMessage { roomId :: !Text }
    deriving (Eq, Show)

instance ToJSON ClientMessageBody where
    toJSON = \case
        ClientSdpMessage sdp -> A.object
            [ ("t", "sdp")
            , ("sdp", toJSON sdp)
            ]
        ClientIceMessage ice -> A.object
            [ ("t", "ice")
            , ("ice", toJSON ice)
            ]
        JoinRoomMessage {roomId} -> A.object
            [ ("t", "joinRoom")
            , ("roomId", toJSON roomId)
            ]
instance FromJSON ClientMessageBody where
    parseJSON = A.withObject "ClientMessage" $ \v -> do
        t <- v .: "t"
        case t of
            A.String "sdp" -> ClientSdpMessage <$> v .: "sdp"
            A.String "ice" -> ClientIceMessage <$> v .: "ice"
            A.String "joinRoom" -> JoinRoomMessage <$> v .: "roomId"
            _ -> fail "Invalid tag value"

data Sdp = Sdp
    { type' :: !Text
    , sdp :: !Text
    }
    deriving (Eq, Show, Generic)
instance ToJSON Sdp where
    toJSON Sdp {type', sdp} = A.object
        [ ("type", toJSON type')
        , ("sdp", toJSON sdp)
        ]
instance FromJSON Sdp where
    parseJSON = A.withObject "SdpMessage" $ \v -> Sdp
        <$> v .: "type"
        <*> v .: "sdp"

data JoinRoom = JoinRoom
    { roomId :: !Text
    }
    deriving (Eq, Show, Generic)
    deriving anyclass (ToJSON, FromJSON)

data ServerMessage = ServerMessage
    { senderId :: !Text
    , body :: !ServerMessageBody
    }
    deriving (Eq, Show, Generic)
    deriving anyclass (ToJSON, FromJSON)
data ServerMessageBody
    = ServerSdpMessage !Sdp
    | ServerIceMessage !A.Value
    | IdentityMessage { roomId :: !Text }
    | PartyJoinsMessage
    | JoinRoomResponse { status :: !Text }
    | LeaveRoomMessage
    deriving (Eq, Show)

instance ToJSON ServerMessageBody where
    toJSON = \case
        ServerSdpMessage sdp -> A.object
            [ ("t", "sdp")
            , ("sdp", toJSON sdp)
            ]
        ServerIceMessage ice -> A.object
            [ ("t", "ice")
            , ("ice", toJSON ice)
            ]
        IdentityMessage { roomId } -> A.object
            [ ("t", "identity")
            , ("roomId", toJSON roomId)
            ]
        PartyJoinsMessage -> A.object [("t", "partyJoins")]
        JoinRoomResponse { status } -> A.object
            [ ("t", "joinRoom")
            , ("status", toJSON status)
            ]
        LeaveRoomMessage -> A.object [("t", "leaveRoom")]
instance FromJSON ServerMessageBody where
    parseJSON = A.withObject "ServerMessage" $ \v -> do
        t <- v .: "t"
        case t of
            A.String "sdp" -> ServerSdpMessage <$> v .: "sdp"
            A.String "ice" -> ServerIceMessage <$> v .: "ice"
            A.String "identity" -> IdentityMessage <$> v .: "roomId"
            A.String "partyJoins" -> pure PartyJoinsMessage
            A.String "joinRoom" -> JoinRoomResponse <$> v .: "status"
            A.String "leaveRoom" -> pure LeaveRoomMessage
            _ -> fail "Invalid tag value"

main :: HasCallStack => IO ()
main = do
    putStrLn "Test: basic join room"
    testJoinRoom
    putStrLn "Test: host gets notified when watcher leaves"
    testHostNotifiedOnWatcherLeave
    putStrLn "Test: multiple watchers"
    testMultipleWatchers
    putStrLn "Test: multiple rooms"
    testMultipleRooms

testJoinRoom :: HasCallStack => IO ()
testJoinRoom = WS.runClient "localhost" 31337 "/" $ \c1 -> WS.runClient "localhost" 31337 "/" $ \c2 -> do
    (hostId, roomId) <- firstMessage c1
    (watcherId, watcherRoomId) <- firstMessage c2
    assertNe roomId watcherRoomId

    -- watcher joins the room
    send c2 ClientMessage
        { targetId = Nothing
        , body = JoinRoomMessage roomId
        }
    -- host receives message about other party
    ServerMessage {senderId = watcherId', body = PartyJoinsMessage} <- recv c1
    assertEq watcherId watcherId'
    -- watcher receives success message
    ServerMessage {senderId, body = JoinRoomResponse {status = "success"}} <- recv c2
    assertEq senderId watcherId

    -- exchange sdp messages
    -- host to watcher
    sdp <- pure Sdp
            { type' = "offer"
            , sdp = "take that you worm"
            }
    send c1 ClientMessage
        { targetId = Just watcherId -- host knows ids
        , body = ClientSdpMessage sdp
        }
    ServerMessage {senderId, body} <- recv c2
    assertEq senderId hostId
    assertEq body (ServerSdpMessage sdp)
    -- watcher to host
    sdp <- pure Sdp
            { type' = "answer"
            , sdp = "you needed worthy opponents"
            }
    send c2 ClientMessage
        { targetId = Nothing -- watcher sends to room owner
        , body = ClientSdpMessage sdp
        }
    ServerMessage {senderId, body} <- recv c1
    assertEq senderId watcherId
    assertEq body (ServerSdpMessage sdp)

    -- send several ice messages in a row
    ice1 <- pure $ A.object [("resident", "human")]
    send c1 ClientMessage
        { targetId = Just watcherId
        , body = ClientIceMessage ice1
        }
    ice2 <- pure $ A.object [("pale", "communion")]
    send c1 ClientMessage
        { targetId = Just watcherId
        , body = ClientIceMessage ice2
        }
    -- receive them all
    ServerMessage {senderId, body} <- recv c2
    assertEq senderId hostId
    assertEq body (ServerIceMessage ice1)
    ServerMessage {senderId, body} <- recv c2
    assertEq senderId hostId
    assertEq body (ServerIceMessage ice2)
    -- the other way around
    ice1 <- pure $ A.object [("cold", "days")]
    send c2 ClientMessage
        { targetId = Nothing
        , body = ClientIceMessage ice1
        }
    ice2 <- pure $ A.object [("lost", "ways")]
    send c2 ClientMessage
        { targetId = Nothing
        , body = ClientIceMessage ice2
        }
    -- receive them all
    ServerMessage {senderId, body} <- recv c1
    assertEq senderId watcherId
    assertEq body (ServerIceMessage ice1)
    ServerMessage {senderId, body} <- recv c1
    assertEq senderId watcherId
    assertEq body (ServerIceMessage ice2)

    -- close host connection
    close c1
    ServerMessage {senderId, body = LeaveRoomMessage} <- recv c2
    assertEq senderId hostId

    pure ()

testHostNotifiedOnWatcherLeave :: IO ()
testHostNotifiedOnWatcherLeave = WS.runClient "localhost" 31337 "/" $ \c1 -> WS.runClient "localhost" 31337 "/" $ \c2 -> do
    (hostId, roomId) <- firstMessage c1
    (watcherId, watcherRoomId) <- firstMessage c2

    -- watcher joins the room
    send c2 ClientMessage
        { targetId = Nothing
        , body = JoinRoomMessage roomId
        }
    -- host receives message about other party
    ServerMessage {senderId = watcherId', body = PartyJoinsMessage} <- recv c1
    assertEq watcherId watcherId'
    -- watcher receives success message
    ServerMessage {senderId, body = JoinRoomResponse {status = "success"}} <- recv c2
    assertEq senderId watcherId

    -- close watcher connection
    close c2
    ServerMessage {senderId, body = LeaveRoomMessage} <- recv c1
    assertEq senderId watcherId

    pure ()

testMultipleWatchers :: IO ()
testMultipleWatchers = WS.runClient "localhost" 31337 "/" $ \c1 -> WS.runClient "localhost" 31337 "/" $ \c2 -> WS.runClient "localhost" 31337 "/" $ \c3 -> do
    (watcher1Id, watcher1RoomId) <- firstMessage c1
    (watcher2Id, watcher2RoomId) <- firstMessage c2
    (hostId, roomId) <- firstMessage c3
    assertNe roomId watcher1RoomId
    assertNe roomId watcher2RoomId

    -- watcher1 joins the room
    send c1 ClientMessage
        { targetId = Nothing
        , body = JoinRoomMessage roomId
        }
    -- watcher2 joins the room
    send c2 ClientMessage
        { targetId = Nothing
        , body = JoinRoomMessage roomId
        }
    -- watcher1 receives success message
    ServerMessage {senderId, body = JoinRoomResponse {status = "success"}} <- recv c1
    assertEq senderId watcher1Id
    -- watcher2 receives success message
    ServerMessage {senderId, body = JoinRoomResponse {status = "success"}} <- recv c2
    assertEq senderId watcher2Id
    -- host receives messages about watchers
    ServerMessage {senderId = watcherId, body = PartyJoinsMessage} <- recv c3
    assertEq watcher1Id watcherId
    ServerMessage {senderId = watcherId, body = PartyJoinsMessage} <- recv c3
    assertEq watcher2Id watcherId

    -- host sends different messages to different clients
    ice1 <- pure $ A.object [("dream", "the dead")]
    ice2 <- pure $ A.object [("will's", "song")]
    send c3 ClientMessage
        { targetId = Just watcher1Id
        , body = ClientIceMessage ice1
        }
    send c3 ClientMessage
        { targetId = Just watcher2Id
        , body = ClientIceMessage ice2
        }
    -- clients receive corresponding messages
    ServerMessage {senderId, body} <- recv c1
    assertEq senderId hostId
    assertEq body (ServerIceMessage ice1)
    ServerMessage {senderId, body} <- recv c2
    assertEq senderId hostId
    assertEq body (ServerIceMessage ice2)

    -- both clients send a message to the host
    ice1 <- pure $ A.object [("the hands", "are the hardest")]
    ice2 <- pure $ A.object [("songs", "for no one")]
    send c1 ClientMessage
        { targetId = Nothing
        , body = ClientIceMessage ice1
        }
    send c2 ClientMessage
        { targetId = Nothing
        , body = ClientIceMessage ice2
        }
    -- host receives them both
    ServerMessage {senderId, body} <- recv c3
    assertEq senderId watcher1Id
    assertEq body (ServerIceMessage ice1)
    ServerMessage {senderId, body} <- recv c3
    assertEq senderId watcher2Id
    assertEq body (ServerIceMessage ice2)

    -- host closes the connection
    close c3
    ServerMessage {senderId, body = LeaveRoomMessage} <- recv c1
    assertEq senderId hostId
    ServerMessage {senderId, body = LeaveRoomMessage} <- recv c2
    assertEq senderId hostId

    pure ()

testMultipleRooms :: IO ()
testMultipleRooms = WS.runClient "localhost" 31337 "/" $ \c1h -> WS.runClient "localhost" 31337 "/" $ \c1w ->
    WS.runClient "localhost" 31337 "/" $ \c2h -> WS.runClient "localhost" 31337 "/" $ \c2w -> do
        (host1Id, room1Id) <- firstMessage c1h
        (watcher1Id, watcher1RoomId) <- firstMessage c1w
        (host2Id, room2Id) <- firstMessage c2h
        (watcher2Id, watcher2RoomId) <- firstMessage c2w
        assertNe room1Id room2Id

        -- watchers joins the rooms
        send c1w ClientMessage
            { targetId = Nothing
            , body = JoinRoomMessage room1Id
            }
        send c2w ClientMessage
            { targetId = Nothing
            , body = JoinRoomMessage room2Id
            }
        -- hosts receive messages about other parties
        ServerMessage {senderId = watcherId, body = PartyJoinsMessage} <- recv c1h
        assertEq watcher1Id watcherId
        ServerMessage {senderId = watcherId, body = PartyJoinsMessage} <- recv c2h
        assertEq watcher2Id watcherId
        -- watchers receive success messages
        ServerMessage {senderId, body = JoinRoomResponse {status = "success"}} <- recv c1w
        assertEq senderId watcher1Id
        ServerMessage {senderId, body = JoinRoomResponse {status = "success"}} <- recv c2w
        assertEq senderId watcher2Id

        -- both clients send a message to their host
        ice1 <- pure $ A.object [("love", "conquers all")]
        ice2 <- pure $ A.object [("fill", "my heart")]
        send c1w ClientMessage
            { targetId = Nothing
            , body = ClientIceMessage ice1
            }
        send c2w ClientMessage
            { targetId = Nothing
            , body = ClientIceMessage ice2
            }
        -- both hosts receive them
        ServerMessage {senderId, body} <- recv c2h
        assertEq senderId watcher2Id
        assertEq body (ServerIceMessage ice2)
        ServerMessage {senderId, body} <- recv c1h
        assertEq senderId watcher1Id
        assertEq body (ServerIceMessage ice1)

        -- clients close the connections
        close c2w
        close c1w
        ServerMessage {senderId, body = LeaveRoomMessage} <- recv c1h
        assertEq senderId watcher1Id
        ServerMessage {senderId, body = LeaveRoomMessage} <- recv c2h
        assertEq senderId watcher2Id

        pure ()


firstMessage :: WS.Connection -> IO (Text, Text)
firstMessage conn = do
    ServerMessage { senderId = myId, body = initialBody } <- recv conn
    IdentityMessage { roomId } <- pure initialBody
    pure (myId, roomId)

decodeIo :: (FromJSON a, MonadFail m) => LazyByteString -> m a
decodeIo x = case A.eitherDecode x of
    Left e -> fail e
    Right r -> pure r

recv :: WS.Connection -> IO ServerMessage
recv conn = WS.receiveData conn >>= decodeIo

send :: WS.Connection -> ClientMessage -> IO ()
send conn msg = WS.sendTextData conn $ A.encode msg

close :: WS.Connection -> IO ()
close conn = WS.sendClose conn ("closing" :: LazyByteString)

forkAsync :: IO a -> IO (MVar ())
forkAsync io = do
    promise <- newEmptyMVar
    _threadId <- forkFinally io $ \_ -> putMVar promise ()
    pure promise

await :: MVar () -> IO ()
await = readMVar

assertNe :: (Eq a, Show a, HasCallStack) => a -> a -> IO ()
assertNe a b
    | a /= b  = pure ()
    | otherwise  = error $ "Assertion failed: " <>  show a <> " /= " <> show b
assertEq :: (Eq a, Show a, HasCallStack) => a -> a -> IO ()
assertEq a b
    | a == b  = pure ()
    | otherwise  = error $ "Assertion failed: " <> show a <> " == " <> show b