module HandshakeCache where

import Control.Concurrent.STM
import Control.Monad
import Data.Functor.Identity
import qualified Data.Map.Strict as Map
         ;import Data.Map.Strict (Map)
import Data.Time.Clock.POSIX
import Network.Socket
import Data.Bool

import Crypto.Hash

import Crypto.Tox
import qualified Data.MinMaxPSQ     as MM
         ;import Data.MinMaxPSQ     (MinMaxPSQ')
import DPut
import DebugTag
import Network.Tox.Crypto.Transport (Handshake, HandshakeData (..))
import Network.Tox.DHT.Handlers     (createCookieSTM)
import Network.Tox.DHT.Transport    (Cookie (..), CookieData (..), NodeInfo,
                                     key2id, nodeInfo)
import Network.Tox.Handshake

data HandshakeCache = HandshakeCache
    { -- Note that currently we are storing sent handshakes keyed by the
      -- locally issued cookie nonce.
      hscTable  :: TVar (MinMaxPSQ' Nonce24 POSIXTime (SecretKey,HandshakeData))
    , hscSend   :: SockAddr -> Handshake Encrypted -> IO ()
    , hscCrypto :: TransportCrypto
    , hscPendingCookies :: TVar (Map (PublicKey,PublicKey) ())
    }


newHandshakeCache :: TransportCrypto -> (SockAddr -> Handshake Encrypted -> IO ()) -> IO HandshakeCache
newHandshakeCache crypto send = atomically $ do
    tbl <- newTVar MM.empty
    pcs <- newTVar Map.empty
    return HandshakeCache
        { hscTable  = tbl
        , hscSend   = send
        , hscCrypto = crypto
        , hscPendingCookies = pcs
        }

getSentHandshake :: HandshakeCache
                    -> SecretKey
                    -> SockAddr
                    -> Cookie Identity -- locally issued
                    -> Cookie Encrypted -- remotely issued
                    -> IO (Maybe (SecretKey, HandshakeData))
getSentHandshake hscache me their_addr (Cookie n24 (Identity cd)) ecookie = do
    now <- getPOSIXTime
    io <- atomically $ do
        m <- checkExpiry now . MM.lookup' n24 <$> readTVar (hscTable hscache)
        case m of
            Just s  -> return $ return $ Just s
            Nothing -> do
                let them = longTermKey cd
                case nodeInfo (key2id $ dhtKey cd) their_addr of
                    Left _ -> return $ return Nothing -- Non-internet address.
                    Right their_node -> do
                        (s,hs) <- cacheHandshakeSTM hscache me them their_node ecookie now
                        return $ do
                            dput XNetCrypto $ "getSentHandshake sending new handshake."
                            hscSend hscache their_addr hs
                            return $ Just s
    r <- io
    dput XNetCrypto $ "getSentHandshake me="++show (key2id $ toPublic me)++" their_addr="++show their_addr++" --> " ++ show r
    return r


checkExpiry :: POSIXTime -> Maybe (POSIXTime,r) -> Maybe r
checkExpiry now m = do
    (tm,s) <- m
    guard $ tm + 5 {- seconds -} > now
    return s

hashCookie :: HashAlgorithm a => Cookie Encrypted -> Digest a
hashCookie (Cookie n24 encrypted)
    = hashFinalize $ hashUpdate (hashUpdate hashInit n24) encrypted

cacheHandshakeSTM :: HandshakeCache
                    -> SecretKey        -- ^ my ToxID key
                    -> PublicKey        -- ^ them
                    -> NodeInfo         -- ^ their DHT node
                    -> Cookie Encrypted -- ^ issued to me by them
                    -> POSIXTime        -- ^ current time
                    -> STM ((SecretKey,HandshakeData), Handshake Encrypted)
cacheHandshakeSTM hscache me them their_node ecookie timestamp = do
    newsession <- transportNewKey (hscCrypto hscache)
    freshCookie <- createCookieSTM timestamp (hscCrypto hscache) their_node them
    n24 <- transportNewNonce (hscCrypto hscache)
    let hsdata = HandshakeData
            { baseNonce   = n24
            , sessionKey  = toPublic newsession
            , cookieHash  = hashCookie ecookie
            , otherCookie = freshCookie
            }
    hs <- encodeHandshake timestamp (hscCrypto hscache) me them ecookie hsdata
    let Cookie cnonce _ = freshCookie
    modifyTVar' (hscTable hscache) $ MM.insertTake' 20 cnonce (newsession,hsdata) timestamp
    return ((newsession,hsdata),hs)

cacheHandshake :: HandshakeCache
                    -> SecretKey
                    -> PublicKey
                    -> NodeInfo
                    -> Cookie Encrypted
                    -> IO (Handshake Encrypted)
cacheHandshake hscache me them their_node ecookie = do
    timestamp <- getPOSIXTime
    dput XNetCrypto $ "cacheHandshake " ++ show (key2id them,ecookie)
    atomically $ snd <$> cacheHandshakeSTM hscache me them their_node ecookie timestamp

haveCachedCookie :: HandshakeCache
                    -> PublicKey
                    -> PublicKey
                    -> STM Bool
haveCachedCookie hscache me them = do
    m <- Map.lookup (me,them) <$> readTVar (hscPendingCookies hscache)
    return $ maybe True (const False) m


setPendingCookie :: HandshakeCache
                    -> PublicKey
                    -> PublicKey
                    -> Bool
                    -> STM ()
setPendingCookie hscache me them pending = do
    modifyTVar' (hscPendingCookies hscache) $ Map.alter (const $ bool Nothing (Just ()) pending)
                                                        (me,them)