{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# 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(..)
    , 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  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 Network.Socket

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 Encrypted))
    | DHTDHTRequest PublicKey (Asymm (f DHTRequest))
    | DHTLanDiscovery NodeId

deriving instance ( Show (f (Cookie Encrypted))
                  , Show (f Ping)
                  , Show (f Pong)
                  , Show (f GetNodes)
                  , Show (f SendNodes)
                  , Show (f CookieRequest)
                  , Show (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 Encrypted)), 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 Encrypted))
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, Ord, 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 Encrypted)
 deriving (Eq, Show)

data Cookie (f :: * -> *) = Cookie Nonce24 (f CookieData)

deriving instance Eq   (f CookieData) => Eq   (Cookie f)
deriving instance Ord  (f CookieData) => Ord  (Cookie f)
deriving instance Show (f CookieData) => Show (Cookie f)

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

instance Serialize (Cookie Encrypted) 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.
 deriving Show

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