-- | This module implements the lossless Tox session protocol.
{-# LANGUAGE TupleSections #-}
module Network.Tox.Session
    ( SessionParams(..)
    , SessionKey
    , Session(..)
    , sTheirUserKey
    , sClose
    , handshakeH
    ) where

import Control.Concurrent.STM
import Control.Monad
import Data.Functor.Identity
import Data.Word
import Network.Socket         (SockAddr)

import Crypto.Tox
import Data.PacketBuffer            (PacketInboundEvent (..))
import Data.Tox.Message
import DPut
import Network.Lossless
import Network.QueryResponse
import Network.SessionTransports
import Network.Tox.Crypto.Transport
import Network.Tox.DHT.Transport    (Cookie (..), key2id, longTermKey)
import Network.Tox.Handshake

-- | Alias for 'SecretKey' to document that it is used as the temporary Tox
-- session key corresponding to the 'PublicKey' we sent in the handshake.
type SessionKey = SecretKey

-- | These inputs to 'handshakeH' indicate how to respond to handshakes, how to
-- assign packets to sessions, and what to do with established sessions after
-- they are made lossless by queuing packets and appending sequence numbers.
data SessionParams = SessionParams
    { -- | The database of secret keys necessary to encrypt handshake packets.
      spCrypto           :: TransportCrypto
      -- | This is used to create sessions and dispatch packets to them.
    , spSessions         :: Sessions (CryptoPacket Encrypted)
      -- | This method returns the session information corresponding to the
      -- cookie pair for the remote address.  If no handshake was sent, this
      -- should send one immediately.  It should return 'Nothing' if anything
      -- goes wrong.
    , spGetSentHandshake :: SecretKey -> SockAddr
                                      -> Cookie Identity
                                      -> Cookie Encrypted
                                      -> IO (Maybe (SessionKey, HandshakeData))
      -- | This method is invoked on each new session and is responsible for
      -- launching any threads necessary to keep the session alive.
    , spOnNewSession     :: Session -> IO ()
    }

-- | After a session is established, this information is given to the
-- 'spOnNewSession' callback.
data Session = Session
    { -- | This is the secret user (toxid) key that corresponds to the
      -- local-end of this session.
      sOurKey            :: SecretKey
      -- | The remote address for this session. (Not unique, see 'sSessionID').
    , sTheirAddr         :: SockAddr
      -- | The information we sent in the handshake for this session.
    , sSentHandshake     :: HandshakeData
      -- | The information we received in a handshake for this session.
    , sReceivedHandshake :: Handshake Identity
      -- | This method can be used to trigger packets to be re-sent given a
      -- list of their sequence numbers.  It should be used when the remote end
      -- indicates they lost packets.
    , sResendPackets     :: [Word32] -> IO ()
      -- | This list of sequence numbers should be periodically polled and if
      -- it is not empty, we should request they re-send these packets.  For
      -- convenience, a lower bound for the numbers in the list is also
      -- returned.  Suggested polling interval: a few seconds.
    , sMissingInbound    :: IO ([Word32],Word32)
      -- | A lossless transport for sending and receiving packets in this
      -- session.  It is up to the caller to spawn the await-loop to handle
      -- inbound packets.
    , sTransport         :: Transport String () CryptoMessage
      -- | A unique small integer that identifies this session for as long as
      -- it is established.
    , sSessionID         :: Int
    }

-- | Helper to obtain the remote ToxID key from the locally-issued cookie
-- associated with the session.
sTheirUserKey :: Session -> PublicKey
sTheirUserKey s = longTermKey $ runIdentity cookie
 where
    Cookie _ cookie = handshakeCookie (sReceivedHandshake s)

-- | Helper to close the 'Transport' associated with a session.
sClose :: Session -> IO ()
sClose s = closeTransport (sTransport s)


-- | Call this whenever a new handshake arrives so that a session is
-- negotiated.  It always returns Nothing which makes it convenient to use with
-- 'Network.QueryResponse.addHandler'.
handshakeH :: SessionParams
              -> SockAddr
              -> Handshake Encrypted
              -> IO (Maybe a)
handshakeH sp saddr handshake = do
    decryptHandshake (spCrypto sp) handshake
        >>= either (\err -> return ())
                   (uncurry $ plainHandshakeH sp saddr)
    return Nothing


plainHandshakeH :: SessionParams
                   -> SockAddr
                   -> SecretKey
                   -> Handshake Identity
                   -> IO ()
plainHandshakeH sp saddr skey handshake = do
    let hd = runIdentity $ handshakeData handshake
        prelude = show saddr ++ " --> "
    dput XNetCrypto $ prelude ++ "handshake: " ++ show (otherCookie hd, baseNonce hd)
    sent <- spGetSentHandshake sp skey saddr (handshakeCookie handshake) (otherCookie hd)
    -- TODO: this is always returning sent = Nothing
    dput XNetCrypto $ " <-- (cached) handshake baseNonce" ++ show (fmap (baseNonce . snd) sent)
    forM_ sent $ \(hd_skey,hd_sent) -> do
    sk <- SessionKeys (spCrypto sp)
                      hd_skey
                      (sessionKey hd)
                      <$> atomically (newTVar $ baseNonce hd)
                      <*> atomically (newTVar $ baseNonce hd_sent)
    m <- newSession (spSessions sp) (\() p -> return p) (decryptPacket sk) saddr
    dput XNetCrypto $ prelude ++ "plainHandshakeH: session " ++ maybe "Nothing" (const "Just") m
    forM_ m $ \(sid, t) -> do
        (t2,resend,getMissing)
            <- lossless (\cp a -> return $ fmap (,a) $ checkLossless $ runIdentity $ pktData cp)
                        (\seqno p _ -> do
                            y <- encryptPacket sk $ bookKeeping seqno p
                            return (lossyness (msgID p) == Lossy, y))
                        ()
                        t
        let _ = t :: TransportA String () (CryptoPacket Identity) (CryptoPacket Encrypted)
            _ = t2 :: Transport String () CryptoMessage
        sendMessage t2 () $ OneByte ONLINE
        spOnNewSession sp Session
            { sOurKey            = skey
            , sTheirAddr         = saddr
            , sSentHandshake     = hd_sent
            , sReceivedHandshake = handshake
            , sResendPackets     = resend
            , sMissingInbound    = getMissing
            , sTransport         = t2
            , sSessionID         = sid
            }
    return ()


-- | The per-session nonce and key state maintained by 'decryptPacket' and
-- 'encryptPacket'.
data SessionKeys = SessionKeys
    { skCrypto        :: TransportCrypto -- ^ Cache of shared-secrets.
    , skMe            :: SessionKey      -- ^ My session key
    , skThem          :: PublicKey       -- ^ Their session key
    , skNonceIncoming :: TVar Nonce24    -- ^ +21845 when a threshold is reached.
    , skNonceOutgoing :: TVar Nonce24    -- ^ +1 on every packet
    }

-- | Decrypt an inbound session packet and update the nonce for the next one.
decryptPacket :: SessionKeys -> SockAddr -> CryptoPacket Encrypted -> IO (Maybe (CryptoPacket Identity, ()))
decryptPacket sk saddr (CryptoPacket n16 ciphered) = do
    (n24,δ) <- atomically $ do
        n <- readTVar (skNonceIncoming sk)
        let δ = n16 - nonce24ToWord16 n
        return ( n `addtoNonce24` fromIntegral δ, δ )
    secret <- lookupSharedSecret (skCrypto sk) (skMe sk) (skThem sk) n24
    case decodePlain =<< decrypt secret ciphered of
        Left e  -> return Nothing
        Right x -> do
            when ( δ  > 43690 )
                $ atomically $ writeTVar (skNonceIncoming sk) (n24 `addtoNonce24` 21845)

            do let them = key2id $ skThem sk
                   CryptoData ack seqno _ = x
                   cm = decodeRawCryptoMsg x
               dput XNetCrypto $ unwords [take 8 (show them),"-->",show (msgID cm),show (n24,ack,seqno)]

            return $ Just ( CryptoPacket n16 (pure x), () )

-- | Encrypt an outbound session packet and update the nonce for the next one.
encryptPacket :: SessionKeys -> CryptoData -> IO (CryptoPacket Encrypted)
encryptPacket sk plain = do
    n24 <- atomically $ do
        n24 <- readTVar (skNonceOutgoing sk)
        modifyTVar' (skNonceOutgoing sk) incrementNonce24
        return n24
    secret <- lookupSharedSecret (skCrypto sk) (skMe sk) (skThem sk) n24
    let ciphered = encrypt secret $ encodePlain $ plain

    do let them = key2id $ skThem sk
           CryptoData ack seqno cm = plain
       dput XNetCrypto $ unwords [take 8 (show them),"<--",show (msgID cm),show (n24,ack,seqno)]

    return $ CryptoPacket (nonce24ToWord16 n24) ciphered


-- | Add sequence information to an outbound packet.
--
-- From spec.md:
--
--   Data in the encrypted packets:
--
--    [our recvbuffers buffer_start, (highest packet number handled + 1), (big endian)]
--    [uint32_t packet number if lossless, sendbuffer buffer_end if lossy, (big endian)]
--    [data]
bookKeeping :: SequenceInfo -> CryptoMessage -> CryptoData
bookKeeping (SequenceInfo seqno ack) m = CryptoData
    { bufferStart = ack   :: Word32
    , bufferEnd   = seqno :: Word32
    , bufferData  = m
    }

-- | Classify an inbound packet as lossy or lossless based on its id byte.
checkLossless :: CryptoData -> PacketInboundEvent CryptoMessage
checkLossless cd@CryptoData{ bufferStart = ack
                           , bufferEnd   = no
                           , bufferData  = x   } = tag no x' ack
 where
    x' = decodeRawCryptoMsg cd
    tag = case lossyness (msgID x') of Lossy -> PacketReceivedLossy
                                       _     -> PacketReceived