{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE TupleSections #-}
module Network.Tox.Crypto.Handlers where

import Network.Tox.Crypto.Transport
import Network.Tox.DHT.Transport (Cookie(..),CookieData(..))
import Crypto.Tox
import Control.Concurrent.STM
import Network.Address
import qualified Data.Map.Strict as Map
import Crypto.Hash
import Control.Applicative
import Control.Monad
import Data.Time.Clock.POSIX
import qualified Data.ByteString  as B
import Control.Lens
import Data.Function
import Data.Serialize             as S
import Data.Word
import GHC.Conc (unsafeIOToSTM)

-- util, todo: move to another module
maybeToEither (Just x) = Right x
maybeToEither Nothing  = Left "maybeToEither"

data NetCryptoSessionStatus = Unaccepted | Accepted | Confirmed
    deriving (Eq,Ord,Show,Enum)


type IOHook addr x = addr -> x -> IO (Maybe (x -> x))
type NetCryptoHook = IOHook SockAddr CryptoMessage


data NetCryptoSession = NCrypto { ncState         :: TVar NetCryptoSessionStatus
                                , ncTheirBaseNonce :: TVar Nonce24 -- base nonce + packet number
                                , ncMyPacketNonce :: TVar Nonce24 -- base nonce + packet number
                                , ncHandShake     :: TVar (Maybe (Handshake Encrypted))
                                , ncCookie        :: TVar (Maybe Cookie)
                                , ncTheirSessionPublic :: Maybe PublicKey
                                , ncSessionSecret :: SecretKey
                                , ncSockAddr      :: SockAddr
                                , ncHooks         :: TVar (Map.Map MessageType [NetCryptoHook])
                                }

data NetCryptoSessions = NCSessions { netCryptoSessions :: TVar (Map.Map SockAddr NetCryptoSession)
                                    , transportCrypto :: TransportCrypto
                                    , defaultHooks :: Map.Map MessageType [NetCryptoHook]
                                    }

newSessionsState :: TransportCrypto -> Map.Map MessageType [NetCryptoHook] ->  IO NetCryptoSessions
newSessionsState crypto hooks = do
    x <- atomically $ newTVar Map.empty
    return NCSessions { netCryptoSessions = x
                      , transportCrypto = crypto
                      , defaultHooks = hooks
                      }

data HandshakeParams 
        = HParam
            { hpTheirBaseNonce :: Maybe Nonce24 -- ignore and generate your own
            , hpOtherCookie :: Maybe Cookie
            , hpTheirSessionKeyPublic :: PublicKey
            , hpMySecretKey :: SecretKey
            , hpCookieRemotePubkey :: PublicKey
            , hpCookieRemoteDhtkey :: PublicKey
            }
newHandShakeData :: TransportCrypto -> HandshakeParams -> HandshakeData
newHandShakeData = error "todo"

cryptoNetHandler :: NetCryptoSessions -> SockAddr -> NetCrypto -> IO (Maybe (NetCrypto -> NetCrypto))
cryptoNetHandler sessions addr (NetHandshake (Handshake (Cookie n24 ecookie) nonce24 encrypted)) = do
    -- Handle Handshake Message
    let crypto = transportCrypto sessions
        allsessions = netCryptoSessions sessions
        anyRight xs f = foldr1 (<|>) $ map f xs
    seckeys <- map fst <$> atomically (readTVar (userKeys crypto))
    symkey <- atomically $ transportSymmetric crypto
    now <- getPOSIXTime
    let lr = do -- Either Monad
            (CookieData cookieTime remotePubkey remoteDhtkey) <- (decodePlain =<< decryptSymmetric symkey n24 ecookie)
            (key,HandshakeData { baseNonce, sessionKey, cookieHash, otherCookie })
                <- anyRight seckeys $ \key -> (key,) <$> (decodePlain =<< decrypt (computeSharedSecret key remotePubkey nonce24) encrypted)
            -- check cookie time < 15 seconds ago
            guard (now - fromIntegral cookieTime < 15)
            -- cookie hash is valid? sha512 of ecookie
            let hinit = hashInit
                hctx = hashUpdate hinit n24
                hctx' = hashUpdate hctx ecookie
                digest = hashFinalize hctx'
            guard (cookieHash == digest)
            -- known friend?
            -- todo
            return 
             HParam
                { hpTheirBaseNonce = Just baseNonce
                , hpOtherCookie    = Just otherCookie
                , hpTheirSessionKeyPublic =  sessionKey
                , hpMySecretKey        = key
                , hpCookieRemotePubkey = remotePubkey
                , hpCookieRemoteDhtkey = remoteDhtkey
                }
    case lr of
        Left _ -> return ()
        Right hp@(HParam
                      { hpTheirBaseNonce = Just theirBaseNonce
                      , hpOtherCookie    = Just otherCookie
                      , hpTheirSessionKeyPublic = theirSessionKey
                      , hpMySecretKey        = key
                      , hpCookieRemotePubkey = remotePublicKey
                      , hpCookieRemoteDhtkey = remoteDhtPublicKey
                      }) -> do
            sessionsmap <- atomically $ readTVar allsessions
            -- Do a lookup, in case we decide to handle the update case differently
            case Map.lookup addr sessionsmap of
                _ -> do -- create new session
                    ncState0 <- atomically $ newTVar Accepted
                    ncTheirBaseNonce0 <- atomically $ newTVar theirBaseNonce
                    n24 <- atomically $ transportNewNonce crypto
                    let myhandshakeData = newHandShakeData crypto hp
                        plain = encodePlain myhandshakeData
                        state = computeSharedSecret key remoteDhtPublicKey n24
                        encrypted = encrypt state plain
                        myhandshake = Handshake { handshakeCookie = otherCookie
                                                , handshakeNonce = n24
                                                , handshakeData = encrypted
                                                }
                    ncMyPacketNonce0 <- atomically $ newTVar (baseNonce myhandshakeData)
                    ncHandShake0 <- atomically $ newTVar (Just myhandshake)
                    cookie0 <- atomically $ newTVar (Just otherCookie)
                    newsession <- generateSecretKey
                    ncHooks0 <- atomically $ newTVar (defaultHooks sessions)
                    let netCryptoSession = 
                            NCrypto { ncState         = ncState0
                                    , ncTheirBaseNonce= ncTheirBaseNonce0
                                    , ncMyPacketNonce   = ncMyPacketNonce0
                                    , ncHandShake     = ncHandShake0
                                    , ncCookie        = cookie0
                                    , ncTheirSessionPublic = Just theirSessionKey
                                    , ncSessionSecret = newsession
                                    , ncSockAddr      = addr
                                    , ncHooks         = ncHooks0
                                    }
                    atomically $ modifyTVar allsessions (Map.insert addr netCryptoSession)
    return Nothing 
cryptoNetHandler sessions addr (NetCrypto (CryptoPacket nonce16 encrypted)) = do
    let crypto = transportCrypto sessions
        allsessions = netCryptoSessions sessions
    sessionsmap <- atomically $ readTVar allsessions
    -- Handle Encrypted Message
    case Map.lookup addr sessionsmap of
        Nothing -> return Nothing -- drop packet, we have no session
        Just (NCrypto {ncState, ncHooks,ncSessionSecret,ncTheirSessionPublic,ncTheirBaseNonce}) -> do
            theirBaseNonce <- atomically $ readTVar ncTheirBaseNonce
            -- Try to decrypt message
            let diff :: Word16
                diff = nonce16 - fromIntegral (last2Bytes theirBaseNonce) -- truncating to Word16
            tempNonce <- addtoNonce24 theirBaseNonce (fromIntegral diff) -- expanding to Word
            let lr = do -- Either Monad --
                        pubkey <- maybeToEither ncTheirSessionPublic
                        decodePlain =<< decrypt (computeSharedSecret ncSessionSecret pubkey tempNonce) encrypted
            case lr of
                Left _ -> return Nothing -- decryption failed, ignore packet
                Right (CryptoData {bufferStart, bufferEnd, bufferData=cm}) -> do -- decryption succeeded,
                    -- TODO: Why do I need bufferStart & bufferEnd?
                    --
                    --   buffer_start = highest packet number handled + 1
                    --                , recvbuffers buffer_start
                    --
                    --   bufferEnd = sendbuffer buffer_end if lossy, otherwise packet number
                    -- update ncTheirBaseNonce if necessary
                    when (diff > 2 * dATA_NUM_THRESHOLD)$
                        atomically $ do
                            y <- readTVar ncTheirBaseNonce
                            -- all because Storable forces IO...
                            x <- unsafeIOToSTM $ addtoNonce24 y (fromIntegral dATA_NUM_THRESHOLD)
                            writeTVar ncTheirBaseNonce y
                    -- then set session confirmed,
                    atomically $ writeTVar ncState Confirmed
                    hookmap <- atomically $ readTVar ncHooks
                    -- if lossy, just run hook
                    if lossyness (msgID cm) == Lossy
                      then
                            case Map.lookup (cm ^. messageType) hookmap of
                                Nothing -> return Nothing -- discarding, because no hooks
                                Just hooks -> flip fix (hooks,cm) $ \loop (hooks,msg) -> do
                                                let _ = cm :: CryptoMessage
                                                case (hooks,cm) of
                                                    ([],_) ->  return Nothing
                                                    (hook:more,cm) -> do
                                                        r <- hook addr cm :: IO (Maybe (CryptoMessage -> CryptoMessage))
                                                        case r of
                                                            Just f -> loop (more,f cm)
                                                            Nothing -> return Nothing -- message consumed
                      else -- Lossless message, so try to restore sequence
                        error "todo try to restore sequence of lossless messages"
    where
        last2Bytes :: Nonce24 -> Word
        last2Bytes (Nonce24 bs) = case S.decode (B.drop 22 bs) of
                                    Right n -> n
                                    _ -> error "unreachable-last2Bytes"
        dATA_NUM_THRESHOLD = 21845 -- = 65535 / 3 


cryptoDefaultHooks = Map.empty