module Main (main) where

import Control.Concurrent.Async (race_)
import Control.Concurrent.STM (atomically, readTBQueue, writeTBQueue)
import Control.Monad (forever, replicateM_)
import Data.ByteString qualified as BS
import Data.ByteString.Lazy qualified as BSL
import Data.IORef (newIORef, atomicModifyIORef')
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 ""

serverRaw :: IO ()
serverRaw = QUIC.runServer [("127.0.0.1", 14443)] \conn stream -> do
  putStrLn "Server accepted connection:"
  QUIC.getConnectionInfo conn >>= print

  query <- QUIC.recvStream stream 4096
  putStrLn $ "Server got query: " <> show query
  QUIC.sendStream stream $ "got yer bytes: " <> query
  _ <- QUIC.recvStream stream 4096
  putStrLn "Server quits"

clientRaw :: IO ()
clientRaw = QUIC.runClient "127.0.0.1" "14443" \conn stream -> do
  putStrLn "Client connected:"
  QUIC.getConnectionInfo conn >>= print

  QUIC.sendStream stream "hi there"
  reply <- QUIC.recvStream stream 4096
  putStrLn $ "Client got reply: " <> show reply
  QUIC.closeStream stream
  putStrLn "Client quits"

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

clientBox :: IO ()
clientBox = do
  QUIC.runClient "127.0.0.1" "14443" \_conn stream -> do
    (writeQ, readQ) <- dummyCodec stream
    atomically $ writeTBQueue writeQ "hi there"
    reply <- atomically $ readTBQueue readQ
    putStrLn $ "Client got reply: " <> show reply

dummyCodec :: QUIC.Stream -> IO (MessageQueues BSL.ByteString BS.ByteString)
dummyCodec = streamCodec id (\chunk -> pure ("", Just chunk))

data ClientMessage
  = Hello
  | Bye
  deriving (Eq, Show, Ord, Generic)

instance QUIC.Serialise ClientMessage

data ServerMessage
  = Ok Int
  deriving (Eq, Show, Ord, Generic)

instance QUIC.Serialise ServerMessage

serverSerialise :: IO ()
serverSerialise = QUIC.runServer [("127.0.0.1", 14443)] \_conn stream -> do
  putStrLn "Server accepted connection:"
  (writeQ, readQ) <- streamSerialise stream
  let
    loop counter = do
      query <- atomically (readTBQueue readQ)
      putStrLn $ "Server got query: " <> show query
      case query of
        Hello -> do
          atomically $ writeTBQueue writeQ (Ok counter)
          loop (counter + 1)
        Bye ->
          pure ()
  loop 0

clientSerialise :: IO ()
clientSerialise = do
  QUIC.runClient "127.0.0.1" "14443" \_conn stream -> do
    (writeQ, readQ) <- streamSerialise stream
    replicateM_ 5 do
      atomically $ writeTBQueue writeQ Hello
      reply <- atomically $ readTBQueue @ServerMessage readQ
      putStrLn $ "Client got reply: " <> show reply
    atomically $ writeTBQueue writeQ Bye

serverSimple :: IO ()
serverSimple = do
  counter <- newIORef 0
  QUIC.runServerSimple "127.0.0.1" 14443 \case
    Hello -> do
      putStrLn "Server got Hello"
      n <- atomicModifyIORef' counter \old -> (old + 1, old)
      pure $ Ok n
    Bye -> do
      putStrLn "Server got Bye"
      error "Whelp, the serverSimple must reply, but the protocol must stop. Needs a re-design."

clientSimple :: IO ()
clientSimple = do
  (stop, call) <- QUIC.startClientSimple "127.0.0.1" "14443"
  replicateM_ 5 do
    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"
  stop