{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE StrictData #-}

module Network.QUIC.Connection.Types where

import Control.Concurrent
import Control.Concurrent.STM
import qualified Crypto.Token as CT
import Data.Array.IO
import Data.ByteString.Internal
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.X509 (CertificateChain)
import Foreign.Marshal.Alloc
import Foreign.Ptr (nullPtr)
import Network.Control (
    Rate,
    RxFlow,
    TxFlow,
    newRate,
    newRxFlow,
    newTxFlow,
 )
import Network.Socket (SockAddr, Socket)
import Network.TLS.QUIC
import System.Mem.Weak (Weak)

import Network.QUIC.Config
import Network.QUIC.Connector
import Network.QUIC.Crypto
import Network.QUIC.Imports
import Network.QUIC.Logger
import Network.QUIC.Parameters
import Network.QUIC.Qlog
import Network.QUIC.Recovery
import Network.QUIC.Stream
import Network.QUIC.Types

----------------------------------------------------------------

dummySecrets :: TrafficSecrets a
dummySecrets = (ClientTrafficSecret "", ServerTrafficSecret "")

----------------------------------------------------------------

data RoleInfo
    = ClientInfo
        { clientInitialToken :: Token -- new or retry token
        , resumptionInfo :: ResumptionInfo
        , incompatibleVN :: Bool
        , sockConnected :: Bool
        }
    | ServerInfo
        { tokenManager :: ~CT.TokenManager
        , registerCID :: CID -> Connection -> IO ()
        , unregisterCID :: CID -> IO ()
        , askRetry :: Bool
        , stopServer :: IO ()
        , certChain :: Maybe CertificateChain
        }

defaultClientRoleInfo :: RoleInfo
defaultClientRoleInfo =
    ClientInfo
        { clientInitialToken = emptyToken
        , resumptionInfo = defaultResumptionInfo
        , incompatibleVN = False
        , sockConnected = False
        }

defaultServerRoleInfo :: RoleInfo
defaultServerRoleInfo =
    ServerInfo
        { tokenManager = undefined
        , registerCID = \_ _ -> return ()
        , unregisterCID = \_ -> return ()
        , askRetry = False
        , stopServer = return ()
        , certChain = Nothing
        }

-- cidInfoSRT in CIDInfo is only used in client
-- which accepts stateless reset.
data CIDDB = CIDDB
    { usedCIDInfo :: CIDInfo
    , cidInfos :: IntMap CIDInfo
    , revInfos :: Map CID Int
    , nextSeqNum :: Int -- only for mine (new)
    , retirePriorTo :: Int
    , triggeredByMe :: Bool -- only for peer's
    }
    deriving (Show)

newCIDDB :: CID -> CIDDB
newCIDDB cid =
    CIDDB
        { usedCIDInfo = cidInfo
        , cidInfos = IntMap.singleton 0 cidInfo
        , revInfos = Map.singleton cid 0
        , nextSeqNum = 1
        , retirePriorTo = 1
        , triggeredByMe = False
        }
  where
    cidInfo = newCIDInfo 0 cid $ StatelessResetToken ""

----------------------------------------------------------------

data MigrationState
    = NonMigration
    | MigrationStarted
    | SendChallenge PathInfo PathData
    | RecvResponse

{- FOURMOLU_DISABLE -}
instance Eq MigrationState where
    NonMigration     == NonMigration     = True
    MigrationStarted == MigrationStarted = True
    RecvResponse     == RecvResponse     = True
    _                == _                = False
{- FOURMOLU_ENABLE -}

----------------------------------------------------------------

data Coder = Coder
    { encrypt :: Buffer -> PlainText -> AssDat -> PacketNumber -> IO Int
    , decrypt :: Buffer -> CipherText -> AssDat -> PacketNumber -> IO Int
    , supplement :: Maybe Supplement
    }

initialCoder :: Coder
initialCoder =
    Coder
        { encrypt = \_ _ _ _ -> return (-1)
        , decrypt = \_ _ _ _ -> return (-1)
        , supplement = Nothing
        }

data Coder1RTT = Coder1RTT
    { coder1RTT :: Coder
    , secretN :: TrafficSecrets ApplicationSecret
    }

initialCoder1RTT :: Coder1RTT
initialCoder1RTT =
    Coder1RTT
        { coder1RTT = initialCoder
        , secretN = (ClientTrafficSecret "", ServerTrafficSecret "")
        }

data Protector = Protector
    { setSample :: Buffer -> IO ()
    , getMask :: IO Buffer
    , unprotect :: Sample -> Mask
    }

initialProtector :: Protector
initialProtector =
    Protector
        { setSample = \_ -> return ()
        , getMask = return nullPtr
        , unprotect = \_ -> Mask ""
        }

----------------------------------------------------------------

data Negotiated = Negotiated
    { tlsHandshakeMode :: HandshakeMode13
    , applicationProtocol :: Maybe NegotiatedProtocol
    , applicationSecretInfo :: ApplicationSecretInfo
    }

initialNegotiated :: Negotiated
initialNegotiated =
    Negotiated
        { tlsHandshakeMode = FullHandshake
        , applicationProtocol = Nothing
        , applicationSecretInfo = ApplicationSecretInfo defaultTrafficSecrets
        }

----------------------------------------------------------------

newtype StreamIdBase = StreamIdBase {fromStreamIdBase :: Int}
    deriving (Eq, Show)

data Concurrency = Concurrency
    { currentStream :: StreamId
    , maxStreams :: StreamIdBase
    }
    deriving (Show)

newConcurrency :: Role -> Direction -> Int -> Concurrency
newConcurrency rl dir n = Concurrency{..}
  where
    bidi = dir == Bidirectional
    currentStream
        | rl == Client = if bidi then 0 else 2
        | otherwise = if bidi then 1 else 3
    maxStreams = StreamIdBase n

----------------------------------------------------------------

type Send = Buffer -> Int -> IO ()
type Recv = IO ReceivedPacket

-- For migration, two SockAddr for the peer are contained.
data PeerInfo = PeerInfo
    { currPathInfo :: PathInfo
    , prevPathInfo :: Maybe PathInfo
    }

data PathInfo = PathInfo
    { peerSockAddr :: SockAddr
    , pathBytesTx :: TVar Int -- TVar for anti amplification
    , pathBytesRx :: TVar Int -- TVar for anti amplification
    , addressValidated :: TVar Bool
    }

newPathInfo :: SockAddr -> IO PathInfo
newPathInfo peerSockAddr = do
    pathBytesTx <- newTVarIO 0
    pathBytesRx <- newTVarIO 0
    addressValidated <- newTVarIO False
    return PathInfo{..}

----------------------------------------------------------------

{- FOURMOLU_DISABLE -}
-- | A quic connection to carry multiple streams.
data Connection = Connection
    { connState         :: ConnState
    , -- Actions
      connDebugLog      :: DebugLogger
    -- ^ A logger for debugging.
    , connQLog          :: QLogger
    , connHooks         :: Hooks
    , connSend          :: ~Send -- ~ for testing
    , connRecv          :: ~Recv -- ~ for testing
    -- Manage
    , connRecvQ         :: RecvQ
    , connSocket        :: IORef Socket
    , genStatelessResetToken :: CID -> StatelessResetToken
    , readers           :: IORef (Map Word64 (Weak ThreadId))
    , mainThreadId      :: ThreadId
    , controlRate       :: Rate
    , -- Info
      roleInfo          :: IORef RoleInfo
    , quicVersionInfo   :: IORef VersionInfo
    , origVersionInfo   :: VersionInfo -- chosenVersion is client's ver in Initial
    -- Mine
    , myParameters      :: Parameters
    , myCIDDB           :: IORef CIDDB
    , -- Peer
      peerParameters    :: IORef Parameters
    , peerCIDDB         :: TVar CIDDB
    , peerInfo          :: IORef PeerInfo
    , -- Queues
      inputQ            :: InputQ
    , cryptoQ           :: CryptoQ
    , outputQ           :: OutputQ
    , outputRate        :: Rate
    , shared            :: Shared
    , delayedAckCount   :: IORef Int
    , delayedAckCancel  :: IORef (IO ())
    , -- State
      peerPacketNumber  :: IORef PacketNumber -- for RTT1
    , streamTable       :: IORef StreamTable
    , myStreamId        :: TVar Concurrency -- C:0 S:1
    , myUniStreamId     :: TVar Concurrency -- C:2 S:3
    , peerStreamId      :: IORef Concurrency -- C:1 S:0
    , peerUniStreamId   :: IORef Concurrency -- C:3 S:2
    , flowTx            :: TVar TxFlow
    , flowRx            :: IORef RxFlow
    , migrationState    :: TVar MigrationState
    , sentRetirePriorTo :: IORef Bool
    , minIdleTimeout    :: IORef Microseconds
    , bytesTx           :: IORef Int
    , bytesRx           :: IORef Int
    , -- TLS
      pendingQ          :: Array EncryptionLevel (TVar [ReceivedPacket])
    , ciphers           :: IOArray EncryptionLevel Cipher
    , coders            :: IOArray EncryptionLevel Coder
    , coders1RTT        :: IOArray Bool Coder1RTT
    , protectors        :: IOArray EncryptionLevel Protector
    , currentKeyPhase   :: IORef (Bool, PacketNumber)
    , negotiated        :: IORef Negotiated
    , connMyAuthCIDs    :: IORef AuthCIDs
    , connPeerAuthCIDs  :: IORef AuthCIDs
    , -- Resources
      encodeBuf         :: Buffer
    , encryptRes        :: SizedBuffer
    , decryptBuf        :: Buffer
    , connResources     :: IORef (IO ())
    , -- Recovery
      connLDCC          :: LDCC
    }
{- FOURMOLU_ENABLE -}

instance KeepQlog Connection where
    keepQlog conn = connQLog conn

instance Connector Connection where
    getRole = role . connState
    getEncryptionLevel = readTVarIO . encryptionLevel . connState
    getMaxPacketSize = readIORef . maxPacketSize . connState
    getConnectionState = readTVarIO . connectionState . connState
    getPacketNumber = readIORef . packetNumber . connState
    getAlive = readIORef . connectionAlive . connState

setDead :: Connection -> IO ()
setDead conn = writeIORef (connectionAlive $ connState conn) False

makePendingQ :: IO (Array EncryptionLevel (TVar [ReceivedPacket]))
makePendingQ = do
    q1 <- newTVarIO []
    q2 <- newTVarIO []
    q3 <- newTVarIO []
    let lst = [(RTT0Level, q1), (HandshakeLevel, q2), (RTT1Level, q3)]
        arr = array (RTT0Level, RTT1Level) lst
    return arr

{- FOURMOLU_DISABLE -}
newConnection
    :: Role
    -> Parameters
    -> VersionInfo
    -> AuthCIDs
    -> AuthCIDs
    -> DebugLogger
    -> QLogger
    -> Hooks
    -> IORef Socket
    -> IORef PeerInfo
    -> RecvQ
    -> Send
    -> Recv
    -> (CID -> StatelessResetToken)
    -> IO Connection
newConnection rl myParameters origVersionInfo myAuthCIDs peerAuthCIDs connDebugLog connQLog connHooks connSocket peerInfo connRecvQ ~connSend ~connRecv genStatelessResetToken = do
    connState         <- newConnState rl
    -- Manage
    readers           <- newIORef Map.empty
    mainThreadId      <- myThreadId
    controlRate       <- newRate
    -- Info
    roleInfo          <- newIORef roleinfo
    quicVersionInfo   <- newIORef origVersionInfo
    -- Mine
    myCIDDB           <- newIORef (newCIDDB myCID)
    -- Peer
    peerParameters    <- newIORef baseParameters
    peerCIDDB         <- newTVarIO (newCIDDB peerCID)
    -- Queus
    inputQ            <- newTQueueIO
    cryptoQ           <- newTQueueIO
    outputQ           <- newTQueueIO
    outputRate        <- newRate
    shared            <- newShared
    delayedAckCount   <- newIORef 0
    delayedAckCancel  <- newIORef (return ())
    -- State
    peerPacketNumber  <- newIORef 0
    streamTable       <- newIORef emptyStreamTable
    myStreamId        <- newTVarIO (newConcurrency rl Bidirectional 0)
    myUniStreamId     <- newTVarIO (newConcurrency rl Unidirectional 0)
    peerStreamId      <- newIORef peerConcurrency
    peerUniStreamId   <- newIORef peerUniConcurrency
    flowTx            <- newTVarIO (newTxFlow 0) -- limit is set in Handshake
    flowRx            <- newIORef (newRxFlow $ initialMaxData myParameters)
    migrationState    <- newTVarIO NonMigration
    sentRetirePriorTo <- newIORef False
    minIdleTimeout    <- newIORef (milliToMicro $ maxIdleTimeout myParameters)
    bytesTx           <- newIORef 0
    bytesRx           <- newIORef 0
    -- TLS
    pendingQ          <- makePendingQ
    ciphers           <- newArray (InitialLevel, RTT1Level) defaultCipher
    coders            <- newArray (InitialLevel, HandshakeLevel) initialCoder
    coders1RTT        <- newArray (False, True) initialCoder1RTT
    protectors        <- newArray (InitialLevel, RTT1Level) initialProtector
    currentKeyPhase   <- newIORef (False, 0)
    negotiated        <- newIORef initialNegotiated
    connMyAuthCIDs    <- newIORef myAuthCIDs
    connPeerAuthCIDs  <- newIORef peerAuthCIDs
    -- Resources
    encodeBuf         <- mallocBytes bufsiz -- used sender or closere
    encryptBuf        <- mallocBytes bufsiz
    let encryptRes = SizedBuffer encryptBuf bufsiz -- used sender
    decryptBuf        <- mallocBytes bufsiz -- used receiver
    connResources     <- newIORef (free encodeBuf >> free encryptBuf >> free decryptBuf)
    -- Recovery
    let put x = atomically $ writeTQueue outputQ $ OutRetrans x
    connLDCC          <- newLDCC connState connQLog put
    return Connection{..}
  where
    isclient = rl == Client
    roleinfo
        | isclient = defaultClientRoleInfo
        | otherwise = defaultServerRoleInfo
    myCID = fromJust $ initSrcCID myAuthCIDs
    peerCID = fromJust $ initSrcCID peerAuthCIDs
    peer
        | isclient = Server
        | otherwise = Client
    peerConcurrency = newConcurrency peer Bidirectional (initialMaxStreamsBidi myParameters)
    peerUniConcurrency = newConcurrency peer Unidirectional (initialMaxStreamsUni myParameters)
    bufsiz = maximumUdpPayloadSize
{- FOURMOLU_ENABLE -}

defaultTrafficSecrets :: (ClientTrafficSecret a, ServerTrafficSecret a)
defaultTrafficSecrets = (ClientTrafficSecret "", ServerTrafficSecret "")

----------------------------------------------------------------

clientConnection
    :: ClientConfig
    -> VersionInfo
    -> AuthCIDs
    -> AuthCIDs
    -> DebugLogger
    -> QLogger
    -> Hooks
    -> IORef Socket
    -> IORef PeerInfo
    -> RecvQ
    -> Send
    -> Recv
    -> (CID -> StatelessResetToken)
    -> IO Connection
clientConnection ClientConfig{..} verInfo myAuthCIDs peerAuthCIDs =
    newConnection Client ccParameters verInfo myAuthCIDs peerAuthCIDs

serverConnection
    :: ServerConfig
    -> VersionInfo
    -> AuthCIDs
    -> AuthCIDs
    -> DebugLogger
    -> QLogger
    -> Hooks
    -> IORef Socket
    -> IORef PeerInfo
    -> RecvQ
    -> Send
    -> Recv
    -> (CID -> StatelessResetToken)
    -> IO Connection
serverConnection ServerConfig{..} verInfo myAuthCIDs peerAuthCIDs =
    newConnection Server scParameters verInfo myAuthCIDs peerAuthCIDs

----------------------------------------------------------------

newtype Input = InpStream Stream deriving (Show)
data Crypto = InpHandshake EncryptionLevel ByteString deriving (Show)

data Output
    = OutControl EncryptionLevel [Frame]
    | OutHandshake [(EncryptionLevel, ByteString)]
    | OutRetrans PlainPacket

type InputQ = TQueue Input
type CryptoQ = TQueue Crypto
type OutputQ = TQueue Output

----------------------------------------------------------------

type SendStreamQ = TQueue TxStreamData

data Shared = Shared
    { sharedCloseSent :: IORef Bool
    , sharedCloseReceived :: IORef Bool
    , shared1RTTReady :: IORef Bool
    , sharedSendStreamQ :: SendStreamQ
    }

newShared :: IO Shared
newShared = do
    sharedCloseSent <- newIORef False
    sharedCloseReceived <- newIORef False
    shared1RTTReady <- newIORef False
    sharedSendStreamQ <- newTQueueIO
    return Shared{..}