{-# OPTIONS_GHC -Wall #-}
module Network.Typed.Socket (
  SSocket,
  SocketFamily(..),
  SocketProtocol(..),
  SockAddr(..),
  SocketStatus(..),
  ShutdownStatus(..),

  socket,
  tcp4Socket,
  tcp6Socket,
  tcpUnixSocket,
  udp4Socket,
  udp6Socket,
  udpUnixSocket,

  connect,
  bind,
  accept,
  listen,
  send,
  sendTo,
  recv,
  recvFrom,
  close,
  shutdownSend,
  shutdownReceive,
  shutdownBoth,

  withSocket,
  withTcp4Socket,
  withTcp6Socket,
  withTcpUnixSocket,
  withUdp4Socket,
  withUdp6Socket,
  withUdpUnixSocket,

  setSocketOption,
  makePortReusable,
  makeAddrReusable
) where

import Prelude (Bool(..), IO, Char, Int, error)
import Data.ByteString (ByteString)

import Control.Applicative
import Control.Exception (bracket)
import Control.Monad
import qualified Network.Socket as NS
import qualified Network.Socket.ByteString as NSB

--------------------------------------------------------------------------------
-- Core Data Types
--------------------------------------------------------------------------------
-- Socket family singletons
data SocketFamily
  = Unix
  | InetV4
  | InetV6

data SocketProtocol
  = Tcp  -- Stream
  | Udp  -- Datagram

data SocketStatus
  = Unconnected
  | Bound
  | Listening
  | Connected
  | Closed

data ShutdownStatus
  = Available
  | CannotReceive
  | CannotSend
  | CannotSendOrReceive

newtype
  SSocket (f :: SocketFamily) (p :: SocketProtocol)
    (s :: SocketStatus) (sh :: ShutdownStatus)
  = SSocket NS.Socket

data SockAddr (f :: SocketFamily) where
  SockAddrInet ::
    !NS.PortNumber -> !NS.HostAddress -> SockAddr 'InetV4
  SockAddrInet6 ::
    NS.PortNumber -> !NS.FlowInfo -> !NS.HostAddress6 -> !NS.ScopeID -> SockAddr 'InetV6
  SockAddrUnix :: ![Char] -> SockAddr 'Unix

--------------------------------------------------------------------------------
-- Core operations
--------------------------------------------------------------------------------

----------------------------------------
-- Socket creation
----------------------------------------
socket :: forall f p.
  (SockFam f, SockProto p) => IO (SSocket f p 'Unconnected 'Available)
socket =
  let x = socketFamily (socketFamily1 :: SocketFamily1 f)
      y = socketProtocol (sockProto1 :: SocketProtocol1 p)
  in SSocket <$> (NS.socket x y NS.defaultProtocol)

tcp4Socket :: IO (SSocket 'InetV4 'Tcp 'Unconnected 'Available)
tcp4Socket = socket

tcp6Socket :: IO (SSocket 'InetV6 'Tcp 'Unconnected 'Available)
tcp6Socket = socket

tcpUnixSocket :: IO (SSocket 'Unix 'Tcp 'Unconnected 'Available)
tcpUnixSocket = socket

udp4Socket :: IO (SSocket 'InetV4 'Udp 'Unconnected 'Available)
udp4Socket = socket

udp6Socket :: IO (SSocket 'InetV6 'Udp 'Unconnected 'Available)
udp6Socket = socket

udpUnixSocket :: IO (SSocket 'Unix 'Udp 'Unconnected 'Available)
udpUnixSocket = socket

-- bracketed, typed helpers
withSocket :: (SockFam f, SockProto p) => (SSocket f p 'Unconnected 'Available -> IO a) -> IO a
withSocket = bracket socket close

withTcp4Socket :: (SSocket 'InetV4 'Tcp 'Unconnected 'Available -> IO a) -> IO a
withTcp4Socket = withSocket

withTcp6Socket :: (SSocket 'InetV6 'Tcp 'Unconnected 'Available -> IO a) -> IO a
withTcp6Socket = withSocket

withTcpUnixSocket :: (SSocket 'Unix 'Tcp 'Unconnected 'Available -> IO a) -> IO a
withTcpUnixSocket = withSocket

withUdp4Socket :: (SSocket 'InetV4 'Udp 'Unconnected 'Available -> IO a) -> IO a
withUdp4Socket = withSocket

withUdp6Socket :: (SSocket 'InetV6 'Udp 'Unconnected 'Available -> IO a) -> IO a
withUdp6Socket = withSocket

withUdpUnixSocket :: (SSocket 'Unix 'Udp 'Unconnected 'Available -> IO a) -> IO a
withUdpUnixSocket = withSocket

----------------------------------------
-- Socket ops
----------------------------------------
connect ::
  SockAddr f -> SSocket f 'Tcp 'Unconnected sh
  -> IO (SSocket f 'Tcp 'Connected 'Available)
connect addr (SSocket s) =
  NS.connect s (toNetworkSockAddr addr) >> return (SSocket s)

bind :: SockAddr f -> SSocket f p 'Unconnected sh -> IO (SSocket f p 'Bound sh)
bind addr (SSocket s) =
  NS.bind s (toNetworkSockAddr addr) >> return (SSocket s)

listen :: Int -> SSocket f 'Tcp 'Bound sh -> IO (SSocket f 'Tcp 'Listening sh)
listen n (SSocket s) = NS.listen s n >> return (SSocket s)

accept :: SSocket f 'Tcp 'Listening sh -> IO (SSocket f 'Tcp 'Connected sh, NS.SockAddr)
accept (SSocket s) = do
  (s', a) <- NS.accept s
  return (SSocket s', a)

----------------------------------------
-- Socket option management
----------------------------------------
setSocketOption :: NS.SocketOption -> Int -> SSocket f p s sh -> IO (SSocket f p s sh)
setSocketOption opt n ss@(SSocket s) = NS.setSocketOption s opt n >> return ss

makePortReusable :: SSocket f p s sh -> IO (SSocket f p s sh)
makePortReusable s = setSocketOption NS.ReusePort 0 s >> return s

makeAddrReusable :: SSocket f p s sh -> IO (SSocket f p s sh)
makeAddrReusable s = setSocketOption NS.ReuseAddr 0 s >> return s

----------------------------------------
-- Socket termination
----------------------------------------
shutdownReceive ::
  CanShutdownReceive sh s ~ 'True =>
  SSocket f p s sh
  -> IO (SSocket f p s (Shutdown sh 'NS.ShutdownReceive))
shutdownReceive (SSocket s) =
  NS.shutdown s NS.ShutdownReceive >> return (SSocket s)

shutdownSend ::
  CanShutdownSend sh s ~ 'True =>
  SSocket f p s sh
  -> IO (SSocket f p s (Shutdown sh 'NS.ShutdownSend))
shutdownSend (SSocket s) =
  NS.shutdown s NS.ShutdownSend >> return (SSocket s)

shutdownBoth ::
  CanShutdownBoth sh s ~ 'True =>
  SSocket f p s sh
  -> IO (SSocket f p s (Shutdown sh 'NS.ShutdownBoth))
shutdownBoth (SSocket s) =
  NS.shutdown s NS.ShutdownBoth >> return (SSocket s)

close :: Closeable s ~ 'True =>
  SSocket f p s sh -> IO (SSocket f p 'Closed sh)
close (SSocket s) = NS.close s >> return (SSocket s)

----------------------------------------
-- Reading and writing
----------------------------------------
send :: CanSend sh ~ 'True
  => ByteString -> SSocket f 'Tcp 'Connected sh -> IO Int
send bs (SSocket s) = NSB.send s bs

recv :: CanReceive sh ~ 'True
  => Int -> SSocket f 'Tcp 'Connected sh -> IO ByteString
recv n (SSocket s) = NSB.recv s n

-- simplified assumptions: require connectionless UDP for sendTo/recvFrom
sendTo :: CanSend sh ~ 'True
  => ByteString -> SockAddr f -> SSocket f 'Udp 'Unconnected sh -> IO Int
sendTo bs addr (SSocket s) =
  NSB.sendTo s bs (toNetworkSockAddr addr)

recvFrom :: (CanReceive sh ~ 'True, SockFam f)
  => Int -> SSocket f 'Udp 'Unconnected sh -> IO (ByteString, SockAddr f)
recvFrom n (SSocket s) = do
  (bs, sa) <- NSB.recvFrom s n
  return (bs, fromNetworkSockAddr sa)

--------------------------------------------------------------------------------
-- Type machinery to make the above possible
--------------------------------------------------------------------------------

-- GADT 1-1
data SocketFamily1 (f :: SocketFamily) where
  SUnix :: SocketFamily1 'Unix
  SInetV4 :: SocketFamily1 'InetV4
  SInetV6 :: SocketFamily1 'InetV6

-- type-level -> term-level
class SockFam (f :: SocketFamily) where
  socketFamily1 :: SocketFamily1 f

instance SockFam 'Unix where socketFamily1 = SUnix
instance SockFam 'InetV4 where socketFamily1 = SInetV4
instance SockFam 'InetV6 where socketFamily1 = SInetV6

-- term-level -> socket API
socketFamily :: SocketFamily1 f -> NS.Family
socketFamily = \case
  SUnix -> NS.AF_UNIX
  SInetV4 -> NS.AF_INET
  SInetV6 -> NS.AF_INET6

data SocketProtocol1 (p :: SocketProtocol) where
  STcp :: SocketProtocol1 'Tcp
  SUdp :: SocketProtocol1 'Udp

class SockProto (p :: SocketProtocol) where
  sockProto1 :: SocketProtocol1 p

instance SockProto 'Tcp where sockProto1 = STcp
instance SockProto 'Udp where sockProto1 = SUdp

socketProtocol :: SocketProtocol1 p -> NS.SocketType
socketProtocol = \case
  STcp -> NS.Stream
  SUdp -> NS.Datagram

type family Closeable (s :: SocketStatus) :: Bool where
  Closeable 'Unconnected = 'True
  Closeable 'Bound = 'True
  Closeable 'Listening = 'True
  Closeable 'Connected = 'True
  Closeable 'Closed = 'False

toNetworkSockAddr :: SockAddr f -> NS.SockAddr
toNetworkSockAddr = \case
  (SockAddrInet p addr) -> NS.SockAddrInet p addr
  (SockAddrInet6 p f addr scope) -> NS.SockAddrInet6 p f addr scope
  (SockAddrUnix s) -> NS.SockAddrUnix s

-- todo: needs a stronger proof
-- leverage (SocketFamily1 f) to guide type-refinement below, but
-- currently fails to establish an isomorphism between NS.SockAddr and
-- SockAddr f
fromNetworkSockAddr :: forall f. SockFam f => NS.SockAddr -> SockAddr f
fromNetworkSockAddr netAddr = case (netAddr, socketFamily1 :: SocketFamily1 f) of
  (NS.SockAddrInet p addr, SInetV4) -> SockAddrInet p addr
  (NS.SockAddrInet6 p f addr scope, SInetV6) -> SockAddrInet6 p f addr scope
  (NS.SockAddrUnix s, SUnix) -> SockAddrUnix s
  _ -> error "impossible"


-- transitions that make sense
type family Shutdown (sh :: ShutdownStatus) (cmd :: NS.ShutdownCmd)
  :: ShutdownStatus where
  Shutdown 'Available 'NS.ShutdownSend = 'CannotSend
  Shutdown 'Available 'NS.ShutdownReceive = 'CannotReceive
  Shutdown 'Available 'NS.ShutdownBoth = 'CannotSendOrReceive
  Shutdown 'CannotSend 'NS.ShutdownReceive = 'CannotSendOrReceive
  Shutdown 'CannotReceive 'NS.ShutdownSend = 'CannotSendOrReceive

type family CanShutdownReceive (sh :: ShutdownStatus) (s :: SocketStatus) :: Bool where
  CanShutdownReceive 'Available 'Connected = 'True
  CanShutdownReceive 'CannotSend 'Connected = 'True
  CanShutdownReceive 'Available 'Listening = 'True
  CanShutdownReceive 'CannotSend 'Listening = 'True

type family CanShutdownSend (sh :: ShutdownStatus) (s :: SocketStatus)  :: Bool where
  CanShutdownSend 'Available 'Connected = 'True
  CanShutdownSend 'CannotSend 'Connected = 'True
  CanShutdownSend 'Available 'Listening = 'True
  CanShutdownSend 'CannotSend 'Listening = 'True

type family CanShutdownBoth (sh :: ShutdownStatus) (s :: SocketStatus)  :: Bool where
  CanShutdownBoth 'Available 'Connected = 'True
  CanShutdownBoth 'Available 'Listening = 'True

type family CanSend (sh :: ShutdownStatus) :: Bool where
  CanSend 'Available = 'True
  CanSend 'CannotReceive = 'True

type family CanReceive (sh :: ShutdownStatus) :: Bool where
  CanReceive 'Available = 'True
  CanReceive 'CannotSend = 'True

-- not really used yet
-- it'd be nice to track these at the type-level and also extend the
-- interface so more socket options can be set, rather than just the
-- Int ones
{-
data SocketOption
  = Debug
  | ReuseAddr
  | SType
  | SoError
  | DontRoute
  | Broadcast
  | SendBuffer
  | RecvBuffer
  | KeepAlive
  | OOBInline
  | TimeToLive
  | MaxSegment
  | NoDelay
  | Cork
  | Linger
  | ReusePort
  | RecvLowWater
  | SendLowWater
  | RecvTimeout
  | SendTimeout
  | UseLoopBack
  | UserTimeout
  | IPv6Only
-}