{-# LANGUAGE DeriveDataTypeable         #-}
{-# LANGUAGE DeriveFoldable             #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE DeriveTraversable          #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE PatternSynonyms            #-}
{-# LANGUAGE ScopedTypeVariables        #-}
module Tox where

import Control.Arrow
import Control.Concurrent.STM
import qualified Crypto.Cipher.Salsa    as Salsa
import qualified Crypto.Cipher.XSalsa   as XSalsa
import Crypto.ECC.Class
import qualified Crypto.Error           as Cryptonite
import Crypto.Error.Types
import qualified Crypto.MAC.Poly1305    as Poly1305
import Crypto.PubKey.Curve25519
import Crypto.PubKey.ECC.Types
import Crypto.Random
import Data.Bool
import Data.ByteArray                   as BA
import Data.ByteString                  (ByteString)
import Data.ByteString                  as B
import qualified Data.ByteString.Base16 as Base16
import qualified Data.ByteString.Char8  as C8
import Data.ByteString.Lazy             (toStrict)
import Data.Data
import Data.IP
import Data.Maybe
import Data.Monoid
import qualified Data.Serialize         as S
import Data.Typeable
import Data.Word
import Foreign.Marshal.Alloc
import Foreign.Ptr
import Foreign.Storable
import GHC.Generics                     (Generic)
import Network.Address                  (Address, fromSockAddr, sockAddrPort,
                                         toSockAddr, withPort)
import Network.QueryResponse
import Network.Socket
import System.Endian

newtype NodeId = NodeId ByteString
 deriving (Eq,Ord,Show,ByteArrayAccess)

instance S.Serialize NodeId where
    get = NodeId <$> S.getBytes 32
    put (NodeId bs) = S.putByteString bs

data NodeInfo = NodeInfo
  { nodeId   :: NodeId
  , nodeIP :: IP
  , nodePort :: PortNumber
  }

nodeAddr :: NodeInfo -> SockAddr
nodeAddr (NodeInfo _ ip port) = toSockAddr ip `withPort` port

nodeInfo :: NodeId -> SockAddr -> Either String NodeInfo
nodeInfo nid saddr
    | Just ip <- fromSockAddr saddr
    , Just port <- sockAddrPort saddr = Right $ NodeInfo nid ip port
    | otherwise                       = Left "Address family not supported."

data TransactionId = TransactionId
 { transactionKey :: Nonce8  -- ^ Used to lookup pending query.
 , cryptoNonce    :: Nonce24 -- ^ Used during the encryption layer.
 }

newtype Method = MessageType Word8
 deriving (Eq, Ord, S.Serialize)

pattern PingType      = MessageType 0
pattern PongType      = MessageType 1
pattern GetNodesType  = MessageType 2
pattern SendNodesType = MessageType 4

instance Show Method where
    showsPrec d PingType        = mappend "PingType"
    showsPrec d PongType        = mappend "PongType"
    showsPrec d GetNodesType    = mappend "GetNodesType"
    showsPrec d SendNodesType   = mappend "SendNodesType"
    showsPrec d (MessageType x) = mappend "MessageType " . showsPrec (d+1) x

newtype Nonce8 = Nonce8 Word64
 deriving (Eq, Ord)

instance ByteArrayAccess Nonce8 where
    length _ = 8
    withByteArray (Nonce8 w64) kont =
        allocaBytes 8 $ \p -> do
            poke (castPtr p :: Ptr Word64) $ toBE64 w64
            kont p

instance Show Nonce8 where
    showsPrec d nonce = quoted (mappend $ bin2hex nonce)

newtype Nonce24 = Nonce24 ByteString
 deriving (Eq, Ord, ByteArrayAccess)

instance Show Nonce24 where
    showsPrec d nonce = quoted (mappend $ bin2hex nonce)

instance S.Serialize Nonce24 where
    get = Nonce24 <$> S.getBytes 24
    put (Nonce24 bs) = S.putByteString bs

quoted :: ShowS -> ShowS
quoted shows s = '"':shows ('"':s)

bin2hex :: ByteArrayAccess bs => bs -> String
bin2hex = C8.unpack . Base16.encode . convert


data Message a = Message
    { msgType    :: Method
    , msgOrigin  :: NodeId
    , msgNonce   :: Nonce24
    , msgPayload :: a
    }
 deriving (Eq, Show, Generic, Functor, Foldable, Traversable)

data Ciphered = Ciphered { cipheredMAC   :: Poly1305.Auth
                         , cipheredBytes :: ByteString }
 deriving Eq

getMessage :: S.Get (Message Ciphered)
getMessage = do
        typ <- S.get
        nid <- S.get
        tid <- S.get
        mac <- Poly1305.Auth . convert <$> S.getBytes 16
        cnt <- S.remaining
        bs <- S.getBytes cnt
        return Message { msgType    = typ
                       , msgOrigin  = nid
                       , msgNonce   = tid
                       , msgPayload = Ciphered mac bs }

putMessage :: Message Ciphered -> S.Put
putMessage (Message {..}) = do
        S.put msgType
        S.put msgOrigin
        S.put msgNonce
        let Ciphered (Poly1305.Auth mac) bs = msgPayload
        S.putByteString (convert mac)
        S.putByteString bs

-- TODO: Cache symmetric keys.
data SecretsCache = SecretsCache
newEmptyCache = return SecretsCache

id2key :: NodeId -> PublicKey
id2key recipient = case publicKey recipient of
    CryptoPassed key -> key
    -- This should never happen because a NodeId is 32 bytes.
    CryptoFailed e   -> error ("Unexpected pattern fail: "++show e)

key2id :: PublicKey -> NodeId
key2id pk = case S.decode (BA.convert pk) of
                Left _ -> error "key2id"
                Right nid -> nid


zeros32 :: Bytes
zeros32 = BA.replicate 32 0

zeros24 :: Bytes
zeros24 = BA.take 24 zeros32

hsalsa20 k n = a <> b
 where
    Salsa.State st = XSalsa.initialize 20 k n
    (_, as) = BA.splitAt 4 st
    (a, xs) = BA.splitAt 16 as
    (_, bs) = BA.splitAt 24 xs
    (b, _ ) = BA.splitAt 16 bs


computeSharedSecret :: SecretKey -> NodeId -> Nonce24 -> (Poly1305.State, XSalsa.State)
computeSharedSecret sk recipient nonce = (hash, crypt)
 where
    -- diffie helman
    shared = ecdh (Proxy :: Proxy Curve_X25519) sk (id2key recipient)
    -- shared secret XSalsa key
    k = hsalsa20 shared zeros24
    -- cipher state
    st0 = XSalsa.initialize 20 k nonce
    -- Poly1305 key
    (rs, crypt) = XSalsa.combine st0 zeros32
    -- Since rs is 32 bytes, this pattern should never fail...
    Cryptonite.CryptoPassed hash = Poly1305.initialize rs


encryptMessage :: SecretKey -> SecretsCache -> NodeId -> Message ByteString -> Message Ciphered
encryptMessage sk _ recipient plaintext
    = withSecret encipherAndHash sk recipient (msgNonce plaintext) <$> plaintext

decryptMessage :: SecretKey -> SecretsCache -> Message Ciphered -> Either String (Message ByteString)
decryptMessage sk _ ciphertext
    = mapM (withSecret decipherAndAuth sk (msgOrigin ciphertext) (msgNonce ciphertext)) ciphertext

withSecret f sk recipient nonce x = f hash crypt x
 where
    (hash, crypt) = computeSharedSecret sk recipient nonce


encipherAndHash :: Poly1305.State -> XSalsa.State -> ByteString -> Ciphered
encipherAndHash hash crypt m = Ciphered a c
  where
    c = fst . XSalsa.combine crypt $ m
    a = Poly1305.finalize . Poly1305.update hash $ c

decipherAndAuth :: Poly1305.State -> XSalsa.State -> Ciphered -> Either String ByteString
decipherAndAuth hash crypt (Ciphered mac c)
    | (a == mac) = Right m
    | otherwise  = Left "decipherAndAuth: auth fail"
  where
    m = fst . XSalsa.combine crypt $ c
    a = Poly1305.finalize . Poly1305.update hash $ c


-- TODO:
-- Represents the encrypted portion of a Tox packet.
-- data Payload a = Payload a !Nonce8
--
-- Generic packet type: Message (Payload ByteString)

parsePacket :: SecretKey -> SecretsCache -> ByteString -> SockAddr -> Either String (Message ByteString, NodeInfo)
parsePacket sk cache bs addr = do ciphered <- S.runGet getMessage bs
                                  msg <- decryptMessage sk cache ciphered
                                  ni <- nodeInfo (msgOrigin msg) addr
                                  return (msg, ni)

encodePacket :: SecretKey -> SecretsCache -> Message ByteString -> NodeInfo -> (ByteString, SockAddr)
encodePacket sk cache msg ni = ( S.runPut . putMessage $ encryptMessage sk cache (nodeId ni) msg
                               , nodeAddr ni )

newClient :: SockAddr -> IO (Client String Method TransactionId NodeInfo (Message ByteString))
newClient addr = do
    udp <- udpTransport addr
    secret <- generateSecretKey
    let pubkey = key2id $ toPublic secret
    cache <- newEmptyCache
    drg <- getSystemDRG
    self <- atomically $ newTVar
             $ NodeInfo pubkey (fromMaybe (toEnum 0) $ fromSockAddr addr)
                               (fromMaybe 0 $ sockAddrPort addr)
    let net = layerTransport (parsePacket secret cache)
                             (encodePacket secret cache)
                             udp
        dispatch tbl = DispatchMethods
            { classifyInbound = classify
            , lookupHandler = handlers
            , tableMethods = tbl
            }
        genNonce24 var (TransactionId nonce8 _) = atomically $ do
            (g,pending) <- readTVar var
            let (bs, g') = randomBytesGenerate 24 g
            writeTVar var (g',pending)
            return $ TransactionId nonce8 (Nonce24 bs)
        client tbl var = Client
            { clientNet           = net
            , clientDispatcher    = dispatch tbl
            , clientErrorReporter = ignoreErrors -- TODO
            , clientPending       = var
            , clientAddress       = atomically (readTVar self)
            , clientResponseId    = genNonce24 var
            }
    if fitsInInt (Proxy :: Proxy Word64)
      then do
        let intmapT = transactionMethods (contramapT intKey intMapMethods) gen
        intmap_var <- atomically $ newTVar (drg, mempty)
        return (client intmapT intmap_var)
      else do
        let mapT = transactionMethods (contramapT nonceKey mapMethods) gen
        map_var <- atomically $ newTVar (drg, mempty)
        return (client mapT map_var)

last8 :: ByteString -> Nonce8
last8 bs
    | let len = B.length bs
    , (len >= 8)
      = Nonce8 $ let bs'     = B.drop (len - 8) bs
                     Right w = S.runGet S.getWord64be bs'
                 in w
    | otherwise
      = Nonce8 0

dropEnd8 :: ByteString -> ByteString
dropEnd8 bs = B.take (B.length bs - 8) bs


classify :: Message ByteString -> MessageClass String Method TransactionId
classify (Message { msgType    = typ
                  , msgPayload = bs
                  , msgNonce   = nonce24 }) = go $ TransactionId (last8 bs) nonce24
 where
    go = case typ of
            PingType      -> IsQuery PingType
            GetNodesType  -> IsQuery GetNodesType
            PongType      -> IsResponse
            SendNodesType -> IsResponse

encodePayload typ (TransactionId (Nonce8 tid) nonce) self dest b
    = Message { msgType    = typ
              , msgOrigin  = nodeId self
              , msgNonce   = nonce
              , msgPayload = S.encode b <> S.runPut (S.putWord64be tid)
              }

decodePayload :: S.Serialize a => Message ByteString -> Either String a
decodePayload msg = S.decode $ dropEnd8 $ msgPayload msg

handler typ f = Just $ MethodHandler decodePayload (encodePayload typ) f

handlers :: Method -> Maybe (MethodHandler String TransactionId NodeInfo (Message ByteString))
handlers PingType     = handler PingType pingH
handlers GetNodesType = error "find_node"
handlers _            = Nothing

data Ping = Ping

pingH :: NodeInfo -> Ping -> IO Ping
pingH = error "pingH"

intKey :: TransactionId -> Int
intKey (TransactionId (Nonce8 w) _) = fromIntegral w

nonceKey :: TransactionId -> Nonce8
nonceKey (TransactionId n _) = n

-- randomBytesGenerate :: ByteArray byteArray => Int -> gen -> (byteArray, gen)
-- gen :: forall gen. DRG gen => gen -> ((Nonce8, Nonce24), gen)
gen :: SystemDRG -> (TransactionId, SystemDRG)
gen g = let (bs, g')  = randomBytesGenerate 24 g
            (ws, g'') = randomBytesGenerate 8 g'
            Right w   = S.runGet S.getWord64be ws
        in ( TransactionId (Nonce8 w) (Nonce24 bs), g'' )