{-# LANGUAGE TupleSections #-}
module Network.Tox.Session where

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

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)
import Network.Tox.Handshake

type SessionKey = SecretKey

data SessionParams = SessionParams
    { spCrypto           :: TransportCrypto
    , spSessions         :: Sessions (CryptoPacket Encrypted)
    , spGetSentHandshake :: SecretKey -> SockAddr
                                      -> Cookie Identity
                                      -> Cookie Encrypted
                                      -> IO (Maybe (SessionKey, HandshakeData))
    , spOnNewSession     :: Session -> IO ()
    }

data Session = Session
    { sOurKey            :: SecretKey
    , sTheirAddr         :: SockAddr
    , sSentHandshake     :: HandshakeData
    , sReceivedHandshake :: Handshake Identity
    , sResendPackets     :: [Word32] -> IO ()
      -- ^ If they request that we re-send certain packets, this method is how
      -- that is accomplished.
    , sMissingInbound    :: IO ([Word32],Word32)
      -- ^ 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.
    , sTransport         :: Transport String () CryptoMessage
    , sSessionID         :: Int
    }

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 $ prelude ++ "plainHandshakeH: cached outgoing: " ++ 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 _ -> encryptPacket sk $ bookKeeping seqno p)
                        ()
                        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 ()

decryptPacket :: SessionKeys -> SockAddr -> CryptoPacket Encrypted -> IO (Maybe (CryptoPacket Identity, ()))
decryptPacket sk saddr (CryptoPacket n16 ciphered) = do
    (n,δ) <- atomically $ do
        n <- readTVar (skNonceIncoming sk)
        let δ = n16 - nonce24ToWord16 n
        return ( n `addtoNonce24` fromIntegral δ, δ )
    secret <- lookupSharedSecret (skCrypto sk) (skMe sk) (skThem sk) n
    case decodePlain =<< decrypt secret ciphered of
        Left e  -> return Nothing
        Right x -> do
            when ( δ  > 43690 )
                $ atomically $ writeTVar (skNonceIncoming sk) (n `addtoNonce24` 21845)
            return $ Just ( CryptoPacket n16 (pure x), () )

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
    return $ CryptoPacket (nonce24ToWord16 n24) ciphered

data SessionKeys = SessionKeys
    { skCrypto        :: TransportCrypto
    , skMe            :: SecretKey
    , skThem          :: PublicKey
    , skNonceIncoming :: TVar Nonce24 -- +21845 when a threshold is reached.
    , skNonceOutgoing :: TVar Nonce24 -- +1 on every 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
    }

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