{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}

-- |
--   Module     : Polysemy.Socket
--   License    : MIT
--   Stability  : experimental
--
-- Socket effect for polysemy.
module Polysemy.Socket
  ( -- * Definition
    Socket (..),
    acquire,
    bind,
    listen,
    send,
    receive,
    release,
    accept,
    close,

    -- * Adapters
    adaptSocket,
    adaptSocketSem,

    -- * Eliminators
    runSocketTcp,
    runSocketVirtual,
    Virtual,

    -- * Tracing
    traceSocket,
  )
where

import Control.Monad
import Data.ByteString (ByteString)
import Polysemy
import Polysemy.Input
import Polysemy.Resource
import Polysemy.State
import Polysemy.Trace
import qualified System.Socket as S
import qualified System.Socket.Family.Inet6 as S
import qualified System.Socket.Protocol.TCP as S
import qualified System.Socket.Type.Stream as S

-- |
--
-- @since 0.0.1.0
data Socket s k b c m a where
  Acquire :: Socket s k b c m s
  Bind :: s -> Socket s k b c m ()
  Release :: s -> Socket s k b c m ()
  Listen :: s -> Int -> Socket s k b c m ()
  Accept :: s -> Socket s k b c m (s, k)
  Send :: s -> c -> Socket s k b c m ()
  Receive :: s -> Socket s k b c m b
  Close :: s -> Socket s k b c m ()

makeSem ''Socket

-- |
--
-- @since 0.0.1.0
traceSocket :: forall s k b c r a. Members '[Socket s k b c, Trace] r => Sem r a -> Sem r a
traceSocket = intercept @(Socket s k b c) $ \case
  Acquire -> trace "acquring" >> acquire @s @k @b @c
  Bind s -> trace "binding" >> bind @s @k @b @c s
  Listen s x -> trace "listening" >> listen @s @k @b @c s x
  Receive s -> trace "receiving" >> receive @s @k @b @c s
  Send s x -> trace "sending" >> send @s @k @b @c s x
  Release s -> trace "releasing" >> release @s @k @b @c s
  Accept s -> trace "accept" >> accept @s @k @b @c s
  Close s -> trace "closing" >> close @s @k @b @c s
{-# INLINE traceSocket #-}

-- |
--
-- @since 0.0.1.0
adaptSocket :: Members '[Socket s k b' c'] r => (b' -> b) -> (c -> c') -> Sem (Socket s k b c ': r) a -> Sem r a
adaptSocket f g = adaptSocketSem (pure . f) (pure . g)
{-# INLINE adaptSocket #-}

-- |
--
-- @since 0.0.1.0
adaptSocketSem :: forall s k b c b' c' r a. Members '[Socket s k b' c'] r => (b' -> Sem r b) -> (c -> Sem r c') -> Sem (Socket s k b c ': r) a -> Sem r a
adaptSocketSem f g = interpret $ \case
  Acquire -> acquire @s @k @b' @c'
  Bind s -> bind @s @k @b' @c' s
  Release s -> release @s @k @b' @c' s
  Listen s x -> listen @s @k @b' @c' s x
  Accept s -> accept @s @k @b' @c' s
  Send s c -> g c >>= send @s @k @b' @c' s
  Receive s -> receive @s @k @b' @c' s >>= f
  Close s -> close @s @k @b' @c' s
{-# INLINE adaptSocketSem #-}

-- |
--
-- @since 0.0.1.0
runSocketTcp ::
  forall s r a.
  ( S.Socket S.Inet6 S.Stream S.TCP ~ s,
    Members '[Resource, Embed IO, Input S.Inet6Port] r
  ) =>
  Sem (Socket s (S.SocketAddress S.Inet6) ByteString ByteString ': r) a ->
  Sem r a
runSocketTcp = interpret $ \case
  Acquire -> embed $ S.socket
  Bind s -> do
    x <- input @S.Inet6Port
    Polysemy.embed $ do
      S.setSocketOption s (S.ReuseAddress True)
      S.setSocketOption s (S.V6Only False)
      S.bind s (S.SocketAddressInet6 S.inet6Any x 0 0)
  Listen s x -> embed $ S.listen s x
  Receive s -> embed $ S.receive s 1024 mempty
  Send s x -> void $ embed $ S.sendAll s x S.msgNoSignal
  Release s -> embed $ S.close s
  Accept s -> embed $ S.accept s
  Close s -> embed $ S.close s
{-# INLINE runSocketTcp #-}

data Virtual = Virtual

-- |
--
-- @since 0.0.1.0
runSocketVirtual :: Members '[State p] r => (c -> p) -> (p -> b) -> Sem (Socket Virtual Virtual b c ': r) a -> Sem r a
runSocketVirtual f g = interpret $ \case
  Acquire -> pure Virtual
  Bind _ -> pure ()
  Release _ -> pure ()
  Listen _ _ -> pure ()
  Accept _ -> pure (Virtual, Virtual)
  Send _ c -> put (f c)
  Receive _ -> g <$> get
  Close _ -> pure ()
{-# INLINE runSocketVirtual #-}