{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE TupleSections              #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE UndecidableInstances #-}
module DHTTransport
    ( parseDHTAddr
    , encodeDHTAddr
    , forwardDHTRequests
    , module ToxAddress
    , DHTMessage(..)
    , Ping(..)
    , Pong(..)
    , GetNodes(..)
    , SendNodes(..)
    , DHTPublicKey
    , CookieRequest
    , Cookie
    , DHTRequest
    , mapMessage
    , encrypt
    , decrypt
    , dhtMessageType
    ) where

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

import Control.Arrow
import Control.Monad
import Data.Bool
import qualified Data.ByteString as B
         ;import Data.ByteString (ByteString)
import Data.Tuple
import Data.Serialize            as S
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 (Assym (f Ping))
    | DHTPong (Assym (f Pong))
    | DHTGetNodes (Assym (f GetNodes))
    | DHTSendNodes (Assym (f SendNodes))
    | DHTCookieRequest (Assym (f CookieRequest))
    | DHTCookie Nonce24 (f Cookie)
    | DHTDHTRequest PublicKey (Assym (f DHTRequest))

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

mapMessage :: forall f b. (forall a. Nonce24 -> f a -> b) -> DHTMessage f -> b
mapMessage f (DHTPing          a) = f (assymNonce a) (assymData a)
mapMessage f (DHTPong          a) = f (assymNonce a) (assymData a)
mapMessage f (DHTGetNodes      a) = f (assymNonce a) (assymData a)
mapMessage f (DHTSendNodes     a) = f (assymNonce a) (assymData a)
mapMessage f (DHTCookieRequest a) = f (assymNonce a) (assymData a)
mapMessage f (DHTDHTRequest  _ a) = f (assymNonce a) (assymData a)
mapMessage f (DHTCookie nonce fcookie) = f nonce fcookie


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

parseDHTAddr :: (ByteString, SockAddr) -> Either (DHTMessage Encrypted8,NodeInfo) (ByteString,SockAddr)
parseDHTAddr (msg,saddr)
    | Just (typ,bs) <- B.uncons msg
    , let right = Right (msg,saddr)
          left  = either (const right) 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 -> left $ fanGet bs getCookie (uncurry DHTCookie) (const $ noReplyAddr saddr)
    0x20 -> left $ fanGet bs getDHTReqest (uncurry DHTDHTRequest) (asymNodeInfo saddr . snd)
    _ -> right

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

dhtMessageType :: DHTMessage Encrypted8 -> ( Word8, Put )
dhtMessageType (DHTPing           a) = (0x00, putAssym a)
dhtMessageType (DHTPong           a) = (0x01, putAssym a)
dhtMessageType (DHTGetNodes       a) = (0x02, putAssym a)
dhtMessageType (DHTSendNodes      a) = (0x04, putAssym a)
dhtMessageType (DHTCookieRequest  a) = (0x18, putAssym a)
dhtMessageType (DHTCookie       n x) = (0x19, put n >> put x)
dhtMessageType (DHTDHTRequest   k a) = (0x20, putPublicKey k >> putAssym a)

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, Assym (Encrypted8 DHTRequest))
getDHTReqest = (,) <$> getPublicKey <*> getAssym

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


getDHT :: Sized a => Get (Assym (Encrypted8 a))
getDHT = getAssym


-- Throws an error if called with a non-internet socket.
direct :: Sized a => ByteString
                  -> SockAddr
                  -> (Assym (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 -> Assym 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
 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)

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
            _ -> 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

-- 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 :: Nonce8    -- no_replay
    , dhtpk      :: PublicKey -- dht public key
    , dhtpkNodes :: SendNodes -- other reachable nodes
    }

-- 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 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

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 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 -> (DHTMessage Encrypted8, NodeInfo)
encrypt crypto msg ni = ( transcode (encryptMessage crypto (id2key $ nodeId ni)) msg
                        , ni )

encryptMessage :: Serialize a =>
                  TransportCrypto ->
                  PublicKey ->
                  Nonce24 -> Either (Nonce8,a) (Assym (Nonce8,a)) -> Encrypted8 a
encryptMessage crypto destKey n (Right assym) = E8 $ ToxCrypto.encrypt secret plain
 where
    secret = computeSharedSecret (transportSecret crypto) destKey n
    plain  = encodePlain $ swap $ assymData assym
encryptMessage crypto destKey n (Left plain) = _todo -- need cached public key.

decrypt :: TransportCrypto -> DHTMessage Encrypted8 -> NodeInfo -> Either String (DHTMessage ((,) Nonce8), NodeInfo)
decrypt crypto msg ni = (, ni) <$> (sequenceMessage $ transcode (decryptMessage crypto) msg)

decryptMessage :: Serialize x =>
                  TransportCrypto
                  -> Nonce24
                  -> Either (Encrypted8 x) (Assym (Encrypted8 x))
                  -> (Either String ∘ ((,) Nonce8)) x
decryptMessage crypto n (Right assymE) = plain8 $ ToxCrypto.decrypt secret e
 where
    secret = computeSharedSecret (transportSecret crypto) (senderKey assymE) n
    E8 e   = assymData assymE
    plain8 = Composed . fmap swap . (>>= decodePlain)
decryptMessage crypto n (Left (E8 e)) = _todo -- need cached public key

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

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