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 qualified Data.Tox.DHT.Multi as Multi
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   :: Multi.SessionAddress -> Handshake Encrypted -> IO Multi.SessionAddress
    , hscCrypto :: TransportCrypto
    , hscPendingCookies :: TVar (Map (PublicKey,PublicKey) ())
    }


newHandshakeCache :: TransportCrypto
                        -> (Multi.SessionAddress -> Handshake Encrypted -> IO Multi.SessionAddress)
                        -> 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
                    -> Multi.SessionAddress
                    -> Cookie Identity -- locally issued
                    -> Cookie Encrypted -- remotely issued
                    -> IO (Maybe (Multi.SessionAddress, (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 (their_addr, s)
            Nothing -> do
                let them = longTermKey cd
                case Multi.nodeInfo (key2id $ dhtKey cd) their_addr of
                    Left e -> return $ do
                        dput XNetCrypto $ "getSentHandshake: " ++ show e
                        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."
                            addr' <- hscSend hscache their_addr hs
                            return $ Just (addr', 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
                    -> Multi.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
                    -> Multi.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

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