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

import ToxAddress
import ToxCrypto hiding (encrypt,decrypt)
import qualified ToxCrypto
import Network.QueryResponse

import Control.Arrow
import Control.Monad
import qualified Data.ByteString as B
         ;import Data.ByteString (ByteString)
import Data.Tuple
import Data.Serialize            as S (Get, Serialize, get, put, runGet)
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))

mapMessage :: forall f b. (forall a. Nonce24 -> f a -> b) -> DHTMessage f -> b
mapMessage f msg = f _todo _todo

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

instance Sized SendNodes where
    size = VarSize $ \(SendNodes ns) -> _nodeFormatSize * length ns

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

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

getDHTReqest :: Get (PublicKey, Assym (Encrypted8 DHTRequest))
getDHTReqest = _todo

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


-- 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 saddr asym = either (error . mappend "asymNodeInfo: ") id $ nodeInfo (NodeId $ 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 saddr = either (error . mappend "noReplyAddr: ") id $ nodeInfo zeroID saddr


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

data DHTRequestPacket = DHTRequestPacket
    { requestTarget :: PublicKey
    , request       :: Assym (Encrypted DHTRequest)
    }

instance Serialize DHTRequestPacket where
    get = _todo
    put = _todo


data DHTRequest
    = NATPing Nonce8
    | NATPong Nonce8
    | DHTPK DHTPublicKey

instance Serialize DHTRequest where
    get = return _todo
    put _ = return () -- todo

-- | 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
    , dhtpk :: PublicKey
    , dhtpkNodes :: SendNodes
    }

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

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

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
newtype CookieResponse = CookieResponse Cookie

data Cookie = Cookie Nonce24 (Encrypted CookieData)

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

instance Serialize Cookie where
    get = return $ Cookie _todo _todo
    put _ = return () -- todo

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

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

instance Serialize CookieRequest where
    get = CookieRequest <$> return _todo
    put (CookieRequest _) = return () -- todo

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) msg, ni)

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

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

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