module Main where

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

import Control.Concurrent.Async (async, cancel, link, race_, replicateConcurrently_)
import Control.Exception (bracket)
import Data.ByteString qualified as BS
import Data.ByteString.Lazy qualified as BSL
import Data.IORef (newIORef, atomicModifyIORef')
import Data.IntMap.Strict qualified as IntMap
import Data.Text (Text)
import Data.Text.Encoding (encodeUtf8, decodeUtf8)
import GHC.Generics (Generic)
import Network.QUIC.Simple qualified as QUIC
import Network.QUIC.Simple.Stream (MessageQueues, streamCodec, streamSerialise)
import System.Timeout (timeout)

main :: IO ()
main = do
  putStrLn "Raw"
  race_ serverRaw clientRaw
  putStrLn ""

  putStrLn "Mailbox"
  race_ serverBox clientBox
  putStrLn ""

  putStrLn "Serialise"
  race_ serverSerialise clientSerialise
  putStrLn ""

  putStrLn "Simple"
  race_ serverSimple clientSimple
  putStrLn ""

  putStrLn "Stateful/Async"
  race_ serverStateful clientAsync
  putStrLn ""

-- * Raw

clientRaw :: IO ()
clientRaw = QUIC.runClient "127.0.0.1" "14443" \conn stream -> do
  -- the initial stream comes pre-requested and ready to go
  putStrLn "Client connected:"
  QUIC.getConnectionInfo conn >>= print

  -- streams have no framing on their own
  -- but that's fine, for now...
  QUIC.sendStream stream "hi there"
  reply <- QUIC.recvStream stream 4096
  putStrLn $ "Client got reply: " <> show reply

  QUIC.closeStream stream
  putStrLn "Client quits"

serverRaw :: IO ()
serverRaw = QUIC.runServer [("127.0.0.1", 14443)] \conn stream -> do
  -- the initial stream comes pre-accepted and ready to go
  putStrLn "Server accepted connection:"
  QUIC.getConnectionInfo conn >>= print

  -- wait until *something* arrives and take it all in
  query <- QUIC.recvStream stream 4096
  putStrLn $ "Server got query: " <> show query
  QUIC.sendStream stream $ "got yer bytes: " <> query

  -- the final bytes will be "", signalling the connection getting closed
  finalBytes <- QUIC.recvStream stream 4096
  putStrLn $ "Server quits after " <> show finalBytes

-- * Stream queues / framing

{- | Stateful codec that splits byte stream into text messages

This allows sending empty messages, but may break if the Text has a NUL in there.
What kind of "text" is that anyway?!
-}
cstringCodec :: QUIC.Stream -> IO (MessageQueues Text Text)
cstringCodec stream = do
  previous <- newIORef BSL.empty -- gotta store partial messages somewhere
  snd <$> streamCodec encode (decode previous) stream
  where
    -- add framing
    encode msg = BSL.fromChunks [encodeUtf8 msg, "\NUL"]

    -- strip framing
    decode previous chunk =
      case BS.break (== 0) chunk of
        ("", "") ->
          -- the connection is closing
          pure ("", Nothing)
        (partial, "") -> do
          -- everything is consumed, but no frame marker is in sight
          atomicModifyIORef' previous \prefix ->
            (prefix <> BSL.fromStrict partial, ())
          pure ("", Nothing) -- don't emit a message just yet
        (suffix, leftovers) -> do
          -- found the marker
          prefix <- atomicModifyIORef' previous ("",)
          pure
            ( BS.drop 1 leftovers -- remove the separator
            , Just . decodeUtf8 . BSL.toStrict $
                prefix <> BSL.fromStrict suffix -- decode the whole message
            )

clientBox :: IO ()
clientBox = do
  QUIC.runClient "127.0.0.1" "14443" \_conn stream -> do
    (writeQ, readQ) <- cstringCodec stream
    -- send all the messages at once
    atomically $ writeTBQueue writeQ "hi there"
    -- with the framing in place, this one is a valid message now
    atomically $ writeTBQueue writeQ ""
    atomically $ writeTBQueue writeQ "and again"
    -- read the replies one by one
    reply1 <- atomically $ readTBQueue readQ
    putStrLn $ "Client got reply 1: " <> show reply1
    reply2 <- atomically $ readTBQueue readQ
    putStrLn $ "Client got reply 2: " <> show reply2
    reply3 <- atomically $ readTBQueue readQ
    putStrLn $ "Client got reply 3: " <> show reply3

serverBox :: IO ()
serverBox = QUIC.runServer [("127.0.0.1", 14443)] \_conn stream -> do
  putStrLn "Server accepted connection:"
  (writeQ, readQ) <- cstringCodec stream
  forever do
    -- a linearised logging echo server
    query <- atomically $ readTBQueue readQ
    putStrLn $ "Server got query: " <> show query
    atomically $ writeTBQueue writeQ $ "got yer bytes: " <> query

-- * Serialised messages

data ClientMessage
  = Hello
  | Bye
  deriving (Eq, Show, Ord, Generic) -- derive Generic so Serialise can do its thing

-- using StandAloneDeriving
instance QUIC.Serialise ClientMessage

data ServerMessage
  = Ok Int
  deriving (Eq, Show, Ord, Generic, QUIC.Serialise) -- using DeriveAnyClass

{- XXX: No turn structure is imposed by the protocol.

The messages get delivered in order, but the client and server
have to coordinate implicitly on when to wait for a reply and
when to go without waiting.

Unless the messages themselves carry call IDs, there's no way to
associate replies with responses besides sending in lockstep.
-}
clientSerialise :: IO ()
clientSerialise = do
  QUIC.runClient "127.0.0.1" "14443" \_conn stream -> do
    (writeQ, readQ) <- snd <$> streamSerialise stream
    replicateM_ 5 do
      -- send messages one by one
      atomically $ writeTBQueue writeQ Hello
      -- wait for the reply before sending another one
      reply <- atomically $ readTBQueue @ServerMessage readQ
      -- here we know that the reply is sent in response to the Hello
      putStrLn $ "Client got reply: " <> show reply
    -- send, but don't wait for the reply
    atomically $ writeTBQueue writeQ Bye
    -- the connection is terminated

serverSerialise :: IO ()
serverSerialise = QUIC.runServer [("127.0.0.1", 14443)] \_conn stream -> do
  putStrLn "Server accepted connection:"
  (writeQ, readQ) <- snd <$> streamSerialise stream
  let
    -- simple state-passing loop
    loop counter = do
      query <- atomically (readTBQueue readQ)
      putStrLn $ "Server got query: " <> show query
      case query of
        Hello -> do
          -- reply at once before reading the next message
          atomically $ writeTBQueue writeQ (Ok counter)
          loop (counter + 1)
        Bye ->
          -- don't reply
          -- XXX: the client will hang indefinitely if waiting for the reply here
          pure ()
  loop 0

-- * Sync client-driven RPC

{- | The wrappers from startClientSimple replace implementation details with a handle.

It also imposes a sync/linearised interaction -- every call gets a response.
-}
clientSimple :: IO ()
clientSimple =
  -- since this is a handle pattern, it should be properly scoped
  bracket open close \(_stop, call) -> do
    -- the flow is the same as in clientSerialise,
    replicateM_ 5 do
      -- no queues, just running a function to get a reply
      Ok n <- call Hello
      putStrLn $ "Client got reply " <> show n
    timeout 1000000 (call Bye) >>= mapM_ \reply ->
      putStrLn $ "Shouldn't happen, the server errors out on this: " <> show reply
    putStrLn "Stopping"
  where
    open = QUIC.startClientSimple "127.0.0.1" "14443"
    close (stop, _call) = stop

{- | Wrappers in runServerSimple ensure that every call gets a response.
-}
serverSimple :: IO ()
serverSimple = do
  counter <- newIORef 0 -- some global state shared by all connections
  QUIC.runServerSimple "127.0.0.1" 14443 \case
    Hello -> do
      -- the connection handler is stateless
      self <- myThreadId -- but its thread id can be used as a key in the global state
      -- not now, though...
      putStrLn $ "Server got Hello from " <> show self

      n <- atomicModifyIORef' counter \old -> (old + 1, old)
      pure $ Ok n
    Bye -> do
      putStrLn "Server got Bye"
      -- trade-offs...
      error "Whelp, the serverSimple handler must reply, but the protocol should stop. Needs a re-design..."

-- * DIY async calls/events

{- | Wrappers from startClientAsync only do the connection setup and provide all the internals.

The actual messages implicitly wrap ClientMessage/ServerMessage with an optional call id.
This allows the messages to be treated as calls (sync request-response) or casts (no-reply messages).

The client now has a state of its own to track the requests in flight.
The calls are still sync, but the client now allows calling from multiple threads
without relying on message ordering.
-}
clientAsync :: IO ()
clientAsync = do
  -- wait for the connection
  (client, _conn, (writeQ, readQ)) <- QUIC.startClientAsync "127.0.0.1" "14443"

  -- the casts require no state, but have no response handlers
  let cast q = atomically $ writeTBQueue writeQ (Nothing, q)

  -- events, like casts have no call id, but originating from server
  -- XXX: this protocol has no notion of server-sent requests though
  events <- newTBQueueIO 16

  calls <- newTVarIO mempty
  void $ async do
    -- a dedicated thread now reads all messages and does the triage
    link client -- exit when the client stops
    forever do
      atomically (readTBQueue readQ) >>= \case
        (Nothing, e) ->
          -- no call id -- this is an event
          atomically $ writeTBQueue events e
        (Just callId, r) ->
          -- call id is present -- this is a response
          atomically $ modifyTVar' calls $ IntMap.insert callId r

  -- issuing calls from multiple threads requires avoiding getting someone else's reply
  counter <- newIORef 0 -- the simplest key is a counter
  let
    call q = do
      callId <- atomicModifyIORef' counter \old -> (old + 1, old)
      atomically $ writeTBQueue writeQ (Just callId, q) -- send the message annotated with call id
      replyVar <- newTVarIO undefined -- XXX: this will not be read until the reply is arrived
      let
        popOrRetry = IntMap.alterF \case
          Just r -> Nothing <$ writeTVar replyVar r -- pop the reply and store it outside, then return
          Nothing -> retry -- the calls map is missing the call id, continue waiting
      atomically $
        -- the caller is blocked until the transaction succeeds
        readTVar calls >>= popOrRetry callId >>= writeTVar calls
      -- the transaction concluded, it is now safe to read the var
      readTVarIO replyVar

  -- send all the requests at once, each in its own thread
  replicateConcurrently_ 5 do
    -- everyone will block until they receive their own response
    Ok n <- call Hello
    putStrLn $ "Client got reply " <> show n

  cast Bye -- no reply is even expected
  cancel client
  -- we have no event handler thread so we read the whole queue
  unread <- atomically (flushTBQueue events)
  -- the server does not send any events though
  putStrLn $ "Client events: " <> show unread

{- | Wrappers in runServerStateful provide a richer interface and impose less.
-}
serverStateful :: IO ()
serverStateful = do
  -- some global state
  conns <- newTVarIO mempty
  connIds <- newIORef 0
  QUIC.runServerStateful "127.0.0.1" 14443 (setup conns connIds) (teardown conns) handler
  where
    -- every connection has a local state too
    setup conns counter _conn writeQ = do
      -- generate an explicit key
      connId <- atomicModifyIORef' counter \old -> (old + 1, old)
      -- don't store ThreadIDs directly!
      me <- myThreadId >>= mkWeakThreadId
      atomically $ modifyTVar' conns $
        -- this per-connection state will be observable from outside
        -- everybody can send messages to this connection without a client to "request" them
        IntMap.insert connId (me, writeQ)
      pure (connId, writeQ, 0 :: Int) -- with writeQ the request handlers can use async replies too

    -- use the global and local state for cleanup ops
    teardown conns _conn (connId, _writeQ, _counter) =
      atomically $ modifyTVar' conns $ IntMap.delete connId

    handler connState@(connId, writeQ, counter) msg = do
      putStrLn $ "Server got " <> show msg
      case msg of
        (Just (rid :: Int), Hello) -> do
          -- this server replies later...
          void $ forkIO do
            threadDelay 1000000
            -- with the writeQ being accessible it is possible to send messages at any time
            atomically $ writeTBQueue writeQ (Just rid, Ok counter)
          -- ... but moves to process the next message ASAP
          let connState' = (connId, writeQ, counter + 1) -- update one part, keep the rest
          pure (connState', Nothing)
        _ ->
          -- ignore casted Hellos and any form of Bye
          pure (connState, Nothing)