{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE TupleSections              #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE UndecidableInstances #-}
module Network.Tox.DHT.Transport
    ( parseDHTAddr
    , encodeDHTAddr
    , forwardDHTRequests
    , module Network.Tox.NodeId
    , DHTMessage(..)
    , Ping(..)
    , Pong(..)
    , GetNodes(..)
    , SendNodes(..)
    , DHTPublicKey(..)
    , FriendRequest(..)
    , NoSpam(..)
    , verifyChecksum
    , CookieRequest(..)
    , CookieResponse(..)
    , Cookie(..)
    , CookieData(..)
    , DHTRequest
    , mapMessage
    , encrypt
    , decrypt
    , dhtMessageType
    , asymNodeInfo
    , putMessage -- Convenient for serializing DHTLanDiscovery
    ) where

import Network.Tox.NodeId
import Crypto.Tox hiding (encrypt,decrypt)
import qualified Crypto.Tox as ToxCrypto
import Network.QueryResponse

import Control.Applicative
import Control.Arrow
import Control.Concurrent.STM
import Control.Monad
import Data.Bool
import qualified Data.ByteString.Char8  as B8
import qualified Data.ByteString  as B
         ;import Data.ByteString  (ByteString)
import Data.Functor.Contravariant
import Data.Maybe
import Data.Monoid
import Data.Serialize             as S
import Data.Tuple
import Data.Word
import Foreign.C (CTime(..))
import Network.Socket
import qualified Data.ByteString.Base64 as Base64
import qualified Data.ByteString.Base16 as Base16
import Data.Char (isSpace)

type DHTTransport = Transport String NodeInfo (DHTMessage Encrypted8)
type HandleHi a = Maybe (Either String (DHTMessage Encrypted8, NodeInfo)) -> IO a


data DHTMessage (f :: * -> *)
    = DHTPing (Asymm (f Ping))
    | DHTPong (Asymm (f Pong))
    | DHTGetNodes (Asymm (f GetNodes))
    | DHTSendNodes (Asymm (f SendNodes))
    | DHTCookieRequest (Asymm (f CookieRequest))
    | DHTCookie Nonce24 (f Cookie)
    | DHTDHTRequest PublicKey (Asymm (f DHTRequest))
    | DHTLanDiscovery NodeId

deriving instance ( Show (f Cookie)
                  , Show (Asymm (f Ping))
                  , Show (Asymm (f Pong))
                  , Show (Asymm (f GetNodes))
                  , Show (Asymm (f SendNodes))
                  , Show (Asymm (f CookieRequest))
                  , Show (Asymm (f DHTRequest))
                  ) => Show (DHTMessage f)

mapMessage :: forall f b. (forall a. Nonce24 -> f a -> b) -> DHTMessage f -> Maybe b
mapMessage f (DHTPing          a)      = Just $ f (asymmNonce a) (asymmData a)
mapMessage f (DHTPong          a)      = Just $ f (asymmNonce a) (asymmData a)
mapMessage f (DHTGetNodes      a)      = Just $ f (asymmNonce a) (asymmData a)
mapMessage f (DHTSendNodes     a)      = Just $ f (asymmNonce a) (asymmData a)
mapMessage f (DHTCookieRequest a)      = Just $ f (asymmNonce a) (asymmData a)
mapMessage f (DHTDHTRequest  _ a)      = Just $ f (asymmNonce a) (asymmData a)
mapMessage f (DHTCookie nonce fcookie) = Just $ f nonce fcookie
mapMessage f (DHTLanDiscovery nid)     = Nothing


instance Sized Ping where size = ConstSize 1
instance Sized Pong where size = ConstSize 1

parseDHTAddr :: TransportCrypto -> (ByteString, SockAddr) -> IO (Either (DHTMessage Encrypted8,NodeInfo) (ByteString,SockAddr))
parseDHTAddr crypto (msg,saddr)
    | Just (typ,bs) <- B.uncons msg
    , let right = return $ Right (msg,saddr)
          left  = either (const right) (return . Left)
    = case typ of
    0x00 -> left $ direct bs saddr DHTPing
    0x01 -> left $ direct bs saddr DHTPong
    0x02 -> left $ direct bs saddr DHTGetNodes
    0x04 -> left $ direct bs saddr DHTSendNodes
    0x18 -> left $ direct bs saddr DHTCookieRequest
    0x19 -> do
        cs <- atomically $ readTVar (pendingCookies crypto)
        let ni = fromMaybe (noReplyAddr saddr) $ do
                    (cnt,key) <- lookup saddr cs <|> listToMaybe (map snd cs)
                    either (const Nothing) Just $ nodeInfo (key2id key) saddr
        left $ fanGet bs getCookie (uncurry DHTCookie) (const $ ni)
    0x20 -> left $ fanGet bs getDHTReqest (uncurry DHTDHTRequest) (asymNodeInfo saddr . snd)
    0x21 -> left $ do
               nid <- runGet get bs
               ni <- nodeInfo nid saddr
               return (DHTLanDiscovery nid, ni)
    _ -> right

encodeDHTAddr :: (DHTMessage Encrypted8,NodeInfo) -> IO (ByteString, SockAddr)
encodeDHTAddr (msg,ni) = return (runPut $ putMessage msg, nodeAddr ni)

dhtMessageType :: ( Serialize (f DHTRequest)
                  , Serialize (f Cookie), Serialize (f CookieRequest)
                  , Serialize (f SendNodes), Serialize (f GetNodes)
                  , Serialize (f Pong), Serialize (f Ping)
                  ) => DHTMessage f -> (Word8, Put)
dhtMessageType (DHTPing           a) = (0x00, putAsymm a)
dhtMessageType (DHTPong           a) = (0x01, putAsymm a)
dhtMessageType (DHTGetNodes       a) = (0x02, putAsymm a)
dhtMessageType (DHTSendNodes      a) = (0x04, putAsymm a)
dhtMessageType (DHTCookieRequest  a) = (0x18, putAsymm a)
dhtMessageType (DHTCookie       n x) = (0x19, put n >> put x)
dhtMessageType (DHTDHTRequest   k a) = (0x20, putPublicKey k >> putAsymm a)
dhtMessageType (DHTLanDiscovery nid) = (0x21, put nid)

putMessage :: DHTMessage Encrypted8 -> Put
putMessage msg = case dhtMessageType msg of (typ,p) -> put typ >> p

getCookie :: Get (Nonce24, Encrypted8 Cookie)
getCookie = get

getDHTReqest :: Get (PublicKey, Asymm (Encrypted8 DHTRequest))
getDHTReqest = (,) <$> getPublicKey <*> getAsymm

-- ## DHT Request packets
--
-- | Length | Contents                  |
-- |:-------|:--------------------------|
-- | `1`    | `uint8_t` (0x20)          |
-- | `32`   | receiver's DHT public key |
--   ...      ...


getDHT :: Sized a => Get (Asymm (Encrypted8 a))
getDHT = getAsymm


-- Throws an error if called with a non-internet socket.
direct :: Sized a => ByteString
                  -> SockAddr
                  -> (Asymm (Encrypted8 a) -> DHTMessage Encrypted8)
                  -> Either String (DHTMessage Encrypted8, NodeInfo)
direct bs saddr f = fanGet bs getDHT f (asymNodeInfo saddr)

-- Throws an error if called with a non-internet socket.
asymNodeInfo :: SockAddr -> Asymm a -> NodeInfo
asymNodeInfo saddr asym = either (error . mappend "asymNodeInfo: ") id $ nodeInfo (key2id $ senderKey asym) saddr


fanGet :: ByteString -> Get x -> (x -> a) -> (x -> b) -> Either String (a,b)
fanGet bs getIt f nid = fmap (f &&& nid) $ runGet getIt bs

-- Throws an error if called with a non-internet socket.
noReplyAddr :: SockAddr -> NodeInfo
noReplyAddr saddr = either (error . mappend "noReplyAddr: ") id $ nodeInfo zeroID saddr


data DHTRequest
    -- #### NAT ping request
    --
    --   Length   Contents
    --  :------- :-------------------------
    --   `1`      `uint8_t` (0xfe)
    --   `1`      `uint8_t` (0x00)
    --   `8`      `uint64_t` random number
    = NATPing Nonce8
    -- #### NAT ping response
    --
    --   Length   Contents
    --  :------- :-----------------------------------------------------------------
    --   `1`      `uint8_t` (0xfe)
    --   `1`      `uint8_t` (0x01)
    --   `8`      `uint64_t` random number (the same that was received in request)
    | NATPong Nonce8
    | DHTPK LongTermKeyWrap
    -- From docs/Hardening_docs.txt
    --
    -- All hardening requests must contain exactly 384 bytes of data. (The data sent
    -- must be padded with zeros if it is smaller than that.)
    --
    -- [byte with value: 02 (get nodes test request)][struct Node_format (the node to
    -- test.)][client_id(32 bytes) the id to query the node with.][padding]
    --
    -- packet id: CRYPTO_PACKET_HARDENING (48)
    | Hardening -- TODO
 deriving Show

instance Sized DHTRequest where
    size = VarSize $ \case
        NATPing   _ -> 10
        NATPong   _ -> 10
        DHTPK wrap -> 1{-typ-} + 32{-key-} + 24{-nonce-}
                        + case size of
                            ConstSize n -> n
                            VarSize f -> f (wrapData wrap)
        Hardening -> 1{-typ-} + 384

instance Serialize DHTRequest where
    get = do
        tag <- get
        case tag :: Word8 of
            0xfe -> do
                direction <- get
                bool NATPong NATPing (direction==(0::Word8)) <$> get
            0x9c -> DHTPK <$> get
            0x30 -> pure Hardening -- TODO: CRYPTO_PACKET_HARDENING
            _ -> fail ("unrecognized DHT request: "++show tag)
    put (NATPing n) = put (0xfe00 :: Word16) >> put n
    put (NATPong n) = put (0xfe01 :: Word16) >> put n
    put (DHTPK pk)  = put (0x9c :: Word8) >> put pk
    put (Hardening) = put (0x30 :: Word8) >> putByteString (B.replicate 384 0) -- TODO

-- DHT public key packet:
-- (As Onion data packet?)
--
-- | Length      | Contents                            |
-- |:------------|:------------------------------------|
-- | `1`         | `uint8_t` (0x9c)                    |
-- | `8`         | `uint64_t` `no_replay`              |
-- | `32`        | Our DHT public key                  |
-- | `[39, 204]` | Maximum of 4 nodes in packed format |
data DHTPublicKey = DHTPublicKey
    { dhtpkNonce :: Word64    -- ^ The `no_replay` number is protection if
                              -- someone tries to replay an older packet and
                              -- should be set to an always increasing number.
                              -- It is 8 bytes so you should set a high
                              -- resolution monotonic time as the value.
    , dhtpk      :: PublicKey -- dht public key
    , dhtpkNodes :: SendNodes -- other reachable nodes
    }
 deriving (Eq, Show)


-- int8_t (0x20 sent over onion, 0x12 for sent over net_crypto)
-- [uint32_t nospam][Message (UTF8) 1 to ONION_CLIENT_MAX_DATA_SIZE bytes]
data FriendRequest = FriendRequest
    { friendNoSpam      :: Word32
    , friendRequestText :: ByteString -- UTF8
    }
 deriving (Eq, Show)


-- When sent as a DHT request packet (this is the data sent in the DHT request
-- packet):
--
--   Length     Contents
--  :--------- :-------------------------------
--   `1`        `uint8_t` (0x9c)
--   `32`       Long term public key of sender
--   `24`       Nonce
--   variable   Encrypted payload
data LongTermKeyWrap = LongTermKeyWrap
    { wrapLongTermKey :: PublicKey
    , wrapNonce       :: Nonce24
    , wrapData        :: Encrypted DHTPublicKey
    }
 deriving Show

instance Serialize LongTermKeyWrap where
    get = LongTermKeyWrap <$> getPublicKey <*> get <*> get
    put (LongTermKeyWrap key nonce dta) = putPublicKey key >> put nonce >> put dta


instance Sized DHTPublicKey where
    -- NOTE: 41 bytes includes the 1-byte tag 0x9c in the size.
    -- WARNING: Serialize instance does not include this byte FIXME
    size = VarSize $ \(DHTPublicKey _ _ nodes) -> 41 + case size of
                                                    ConstSize nodes -> nodes
                                                    VarSize sznodes -> sznodes nodes

instance Sized Word32 where size = ConstSize 4

-- FIXME: Inconsitently, this type does not include the 0x20 or 0x12 tag byte
-- where the DHTPublicKey type does include its tag.
instance Sized FriendRequest where
    size = contramap friendNoSpam size <> contramap friendRequestText (VarSize B.length)

instance Serialize DHTPublicKey where
    -- TODO: This should agree with Sized instance.
    get = DHTPublicKey <$> get <*> getPublicKey <*> get
    put (DHTPublicKey nonce key nodes) = do
        put nonce
        putPublicKey key
        put nodes

instance Serialize FriendRequest where
    get = FriendRequest <$> get <*> (remaining >>= getBytes)
    put (FriendRequest nospam txt) = put nospam >> putByteString txt

newtype GetNodes = GetNodes NodeId
    deriving (Eq,Ord,Show,Read,S.Serialize)

instance Sized GetNodes where
    size = ConstSize 32 -- TODO This right?

newtype SendNodes = SendNodes [NodeInfo]
    deriving (Eq,Ord,Show,Read)

instance Sized SendNodes where
    size = VarSize $ \(SendNodes ns) -> case size of
            ConstSize nodeFormatSize -> nodeFormatSize * length ns
            VarSize   nsize          -> sum $ map nsize ns

instance S.Serialize SendNodes where
    get = do
        cnt <- S.get :: S.Get Word8
        ns <- sequence $ replicate (fromIntegral cnt) S.get
        return $ SendNodes ns

    put (SendNodes ns) = do
        let ns' = take 4 ns
        S.put (fromIntegral (length ns') :: Word8)
        mapM_ S.put ns'

data Ping = Ping deriving Show
data Pong = Pong deriving Show

instance S.Serialize Ping where
    get = do w8 <- S.get
             if (w8 :: Word8) /= 0
                then fail "Malformed ping."
                else return Ping
    put Ping = S.put (0 :: Word8)

instance S.Serialize Pong where
    get = do w8 <- S.get
             if (w8 :: Word8) /= 1
                then fail "Malformed pong."
                else return Pong
    put Pong = S.put (1 :: Word8)

newtype CookieRequest = CookieRequest PublicKey
 deriving (Eq, Show)
newtype CookieResponse = CookieResponse Cookie
 deriving (Eq, Show)

data Cookie = Cookie Nonce24 (Encrypted CookieData)
 deriving (Eq, Ord, Show)

instance Sized Cookie where size = ConstSize 112 -- 24 byte nonce + 88 byte cookie data

instance Serialize Cookie where
    get = Cookie <$> get <*> get
    put (Cookie nonce dta) = put nonce >> put dta

data CookieData = CookieData   --   16 (mac)
    { cookieTime  :: Word64    --    8
    , longTermKey :: PublicKey --   32
    , dhtKey      :: PublicKey -- + 32
    }                          -- = 88 bytes when encrypted.

instance Sized CookieData where
    size = ConstSize 72

instance Serialize CookieData where
    get = CookieData <$> get
                     <*> (id2key <$> get)
                     <*> (id2key <$> get)
    put (CookieData tm userkey dhtkey) = do
        put tm
        put (key2id userkey)
        put (key2id dhtkey)

instance Sized CookieRequest where
    size = ConstSize 64 -- 32 byte key + 32 byte padding

instance Serialize CookieRequest where
    get = CookieRequest <$> getPublicKey <* {- padding -} getPublicKey
    put (CookieRequest k) = putPublicKey k >> {- padding -} putPublicKey k

forwardDHTRequests :: TransportCrypto -> (PublicKey -> IO (Maybe NodeInfo)) -> DHTTransport -> DHTTransport
forwardDHTRequests crypto closeLookup dht = dht { awaitMessage = await' }
 where
    await' :: HandleHi a -> IO a
    await' pass = awaitMessage dht $ \case
        Just (Right (m@(DHTDHTRequest target payload),src)) | target /= transportPublic crypto
            -> do mni <- closeLookup target
                  -- Forward the message if the target is in our close list.
                  forM_ mni $ \ni -> sendMessage dht ni m
                  await' pass
        m -> pass m

encrypt :: TransportCrypto -> DHTMessage ((,) Nonce8) -> NodeInfo -> IO (DHTMessage Encrypted8, NodeInfo)
encrypt crypto msg ni = do
    let cipher n plain = Composed $ encryptMessage crypto (id2key $ nodeId ni) n plain
    m <- sequenceMessage $ transcode cipher msg
    return (m, ni)

encryptMessage :: Serialize a =>
                  TransportCrypto ->
                  PublicKey ->
                  Nonce24 -> Either (Nonce8,a) (Asymm (Nonce8,a)) -> IO (Encrypted8 a)
encryptMessage crypto destKey n arg = do
    let plain  = encodePlain $ swap $ either id asymmData arg
    secret <- lookupSharedSecret crypto (transportSecret crypto) destKey n
    return $ E8 $ ToxCrypto.encrypt secret plain

decrypt :: TransportCrypto -> DHTMessage Encrypted8 -> NodeInfo -> IO (Either String (DHTMessage ((,) Nonce8), NodeInfo))
decrypt crypto msg ni = do
    let decipher n c = Composed $ decryptMessage crypto n . left ((,) $ id2key $ nodeId ni) $ c
    msg' <- sequenceMessage $ transcode decipher msg
    return $ fmap (, ni) $ sequenceMessage msg'

decryptMessage :: Serialize x =>
                  TransportCrypto
                  -> Nonce24
                  -> Either (PublicKey, Encrypted8 x) (Asymm (Encrypted8 x))
                  -> IO ((Either String ∘ ((,) Nonce8)) x)
decryptMessage crypto n arg = do
    let (remotekey,E8 e) = either id (senderKey &&& asymmData) arg
        plain8           = Composed . fmap swap . (>>= decodePlain)
    secret <- lookupSharedSecret crypto (transportSecret crypto) remotekey n
    return $ plain8 $ ToxCrypto.decrypt secret e

sequenceMessage :: Applicative m => DHTMessage (m ∘ f) -> m (DHTMessage f)
sequenceMessage (DHTPing asym)              = fmap DHTPing          $ sequenceA $ fmap uncomposed asym
sequenceMessage (DHTPong asym)              = fmap DHTPong          $ sequenceA $ fmap uncomposed asym
sequenceMessage (DHTGetNodes asym)          = fmap DHTGetNodes      $ sequenceA $ fmap uncomposed asym
sequenceMessage (DHTSendNodes asym)         = fmap DHTSendNodes     $ sequenceA $ fmap uncomposed asym
sequenceMessage (DHTCookieRequest asym)     = fmap DHTCookieRequest $ sequenceA $ fmap uncomposed asym
sequenceMessage (DHTCookie n dta)           = fmap (DHTCookie n)    $ uncomposed dta
sequenceMessage (DHTDHTRequest pubkey asym) = fmap (DHTDHTRequest pubkey) $ sequenceA $ fmap uncomposed asym
sequenceMessage (DHTLanDiscovery nid)       = pure $ DHTLanDiscovery nid

transcode :: forall f g. (forall a. Serialize a => Nonce24 -> Either (f a) (Asymm (f a)) -> g a) -> DHTMessage f -> DHTMessage g
transcode f (DHTPing asym)              = DHTPing              $ asym { asymmData = f (asymmNonce asym) (Right asym) }
transcode f (DHTPong asym)              = DHTPong              $ asym { asymmData = f (asymmNonce asym) (Right asym) }
transcode f (DHTGetNodes asym)          = DHTGetNodes          $ asym { asymmData = f (asymmNonce asym) (Right asym) }
transcode f (DHTSendNodes asym)         = DHTSendNodes         $ asym { asymmData = f (asymmNonce asym) (Right asym) }
transcode f (DHTCookieRequest asym)     = DHTCookieRequest     $ asym { asymmData = f (asymmNonce asym) (Right asym) }
transcode f (DHTCookie n dta)           = DHTCookie n          $ f n  $ Left dta
transcode f (DHTDHTRequest pubkey asym) = DHTDHTRequest pubkey $ asym { asymmData = f (asymmNonce asym) (Right asym) }
transcode f (DHTLanDiscovery nid)       = DHTLanDiscovery nid