{-# LANGUAGE DeriveGeneric              #-}
{-# 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
    , toxSpace
    ) where

import qualified Network.Kademlia.Routing as R
import Network.Tox.NodeId
import qualified Network.Tox.TCP.NodeId as TCP
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.Hashable
import Data.Maybe
import Data.Monoid
import Data.Serialize             as S
import Data.Tuple
import Data.Word
import GHC.Generics
import Network.Socket
import DPut
import DebugTag

type DHTTransport ni = Transport String ni (DHTMessage Encrypted8)
type HandleHi ni a = Arrival String ni (DHTMessage Encrypted8) -> 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 :: (Eq saddr, Show ni) =>
                    (saddr -> STM (Maybe ni))
                    -> (NodeId -> saddr -> Either String ni)
                    -> (ByteString, saddr)
                    -> STM (Either (DHTMessage Encrypted8,ni) (ByteString,saddr))
parseDHTAddr pendingCookies nodeInfo (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 nodeInfo bs saddr DHTPing
    0x01 -> left $ direct nodeInfo bs saddr DHTPong
    0x02 -> left $ direct nodeInfo bs saddr DHTGetNodes
    0x04 -> left $ direct nodeInfo bs saddr DHTSendNodes
    0x18 -> left $ direct nodeInfo bs saddr DHTCookieRequest
    0x19 -> do
        mni <- pendingCookies saddr
        let ni = fromMaybe (noReplyAddr nodeInfo saddr) mni
            runio :: IO () -> STM ()
            runio _ = return () -- TODO: run IO action
        runio $ dput XMan $ "Got encrypted cookie! mni="++show mni
        left $ fanGet bs getCookie (uncurry DHTCookie) (const $ ni)
    0x20 -> left $ fanGet bs getDHTReqest (uncurry DHTDHTRequest) (asymNodeInfo nodeInfo saddr . snd)
    0x21 -> left $ do
               nid <- runGet get bs
               ni <- nodeInfo nid saddr
               return (DHTLanDiscovery nid, ni)
    _ -> right

encodeDHTAddr :: (ni -> saddr) -> (DHTMessage Encrypted8,ni) -> IO (ByteString, saddr)
encodeDHTAddr nodeAddr (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 => (NodeId -> saddr -> Either String ni)
                  -> ByteString
                  -> saddr
                  -> (Asymm (Encrypted8 a) -> DHTMessage Encrypted8)
                  -> Either String (DHTMessage Encrypted8, ni)
direct nodeInfo bs saddr f = fanGet bs getDHT f (asymNodeInfo nodeInfo saddr)

-- Throws an error if called with a non-internet socket.
asymNodeInfo :: (NodeId -> saddr -> Either String ni) -> saddr -> Asymm a -> ni
asymNodeInfo nodeInfo 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 :: (NodeId -> saddr -> Either String ni) -> saddr -> ni
noReplyAddr nodeInfo 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
    size = VarSize $ \(DHTPublicKey _ _ nodes) -> 40 + case size of
                                                    ConstSize nodes -> nodes
                                                    VarSize sznodes -> sznodes nodes

instance Sized Word32 where size = ConstSize 4

instance Sized FriendRequest where
    size = contramap friendNoSpam size <> contramap friendRequestText (VarSize B.length)

getTCPNodeList :: S.Get [TCP.NodeInfo]
getTCPNodeList = do
    n <- S.get
    (:) n <$> (getTCPNodeList <|> pure [])

instance Serialize DHTPublicKey where
    -- TODO: This should agree with Sized instance.
    get = DHTPublicKey <$> get <*> getPublicKey <*> (SendNodes <$> getTCPNodeList)
    put (DHTPublicKey nonce key (SendNodes nodes)) = do
        put nonce
        putPublicKey key
        mapM_ 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 [TCP.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)
deriving instance Generic (f CookieData) => Generic (Cookie f)

instance Hashable (Cookie Encrypted)

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, Generic)

instance Sized CookieData where
    size = ConstSize 72

instance Serialize CookieData where
    get = CookieData <$> get <*> getPublicKey <*> getPublicKey
    put (CookieData tm userkey dhtkey) = do
        put tm
        putPublicKey userkey
        putPublicKey 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 ni)) -> DHTTransport ni -> DHTTransport ni
forwardDHTRequests crypto closeLookup dht = dht { awaitMessage = await' }
 where
    -- await' :: HandleHi ni a -> STM (IO a)
    await' = do
        (m, io) <- awaitMessage dht
        return $ case m of
            Arrival src m@(DHTDHTRequest target payload) | target /= transportPublic crypto
                -> (,) Discarded $ do
                      io
                      mni <- closeLookup target
                      -- Forward the message if the target is in our close list.
                      forM_ mni $ \ni -> sendMessage dht ni m
            _ -> (m,io)

encrypt :: TransportCrypto -> (ni -> NodeId) -> DHTMessage ((,) Nonce8) -> ni -> IO (DHTMessage Encrypted8, ni)
encrypt crypto nodeId 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 -> (ni -> NodeId) -> DHTMessage Encrypted8 -> ni -> STM (Either String (DHTMessage ((,) Nonce8), ni))
decrypt crypto nodeId 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))
                  -> STM ((Either String ∘ ((,) Nonce8)) x)
decryptMessage crypto n arg = do
    let (remotekey,E8 e) = either id (senderKey &&& asymmData) arg
        plain8           = Composed . fmap swap . (>>= decodePlain)
    secret <- lookupSharedSecretSTM 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

toxSpace :: R.KademliaSpace NodeId NodeInfo
toxSpace = R.KademliaSpace
    { R.kademliaLocation = nodeId
    , R.kademliaTestBit  = testNodeIdBit
    , R.kademliaXor      = xorNodeId
    , R.kademliaSample   = sampleNodeId
    }