{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DeriveDataTypeable         #-}
{-# LANGUAGE DeriveFoldable             #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE DeriveTraversable          #-}
{-# LANGUAGE ExistentialQuantification  #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE NamedFieldPuns             #-}
{-# LANGUAGE PatternSynonyms            #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TupleSections              #-}
module Network.Tox where

import Debug.Trace
import Control.Exception hiding (Handler)
import Control.Applicative
import Control.Arrow
import Control.Concurrent                     (MVar)
import Control.Concurrent.STM
import Control.Monad
import Control.Monad.Fix
import qualified Crypto.Cipher.ChaChaPoly1305 as Symmetric
import qualified Crypto.Cipher.Salsa          as Salsa
import qualified Crypto.Cipher.XSalsa         as XSalsa
import qualified Crypto.Error                 as Cryptonite
#ifdef CRYPTONITE_BACKPORT
import Crypto.ECC.Class
import Crypto.Error.Types
#else
import Crypto.ECC
import Crypto.Error
#endif
import qualified Crypto.MAC.Poly1305          as Poly1305
import Crypto.PubKey.Curve25519
import Crypto.PubKey.ECC.Types
import Crypto.Random
import qualified Data.Aeson                   as JSON
         ;import Data.Aeson                   (FromJSON, ToJSON, (.=))
import Data.Bitraversable                     (bisequence)
import Data.Bits
import Data.Bits.ByteString                   ()
import Data.Bool
import qualified Data.ByteArray               as BA
         ;import Data.ByteArray               (ByteArrayAccess, Bytes)
import qualified Data.ByteString              as B
         ;import Data.ByteString              (ByteString)
import qualified Data.ByteString.Base16       as Base16
import qualified Data.ByteString.Char8        as C8
import Data.ByteString.Lazy                   (toStrict)
import Data.Char
import Data.Conduit (Source,Sink,Flush(..))
import Data.Data
import Data.Functor.Contravariant
import Data.Hashable
import Data.IP
import Data.Maybe
import qualified Data.MinMaxPSQ               as MinMaxPSQ
         ;import Data.MinMaxPSQ               (MinMaxPSQ')
import Data.Monoid
import Data.Ord
import qualified Data.Serialize               as S
import Data.Time.Clock.POSIX                  (POSIXTime, getPOSIXTime)
import Data.Typeable
import Data.Word
import qualified Data.Wrapper.PSQ             as PSQ
         ;import Data.Wrapper.PSQ             (PSQ)
import qualified Data.Wrapper.PSQInt          as Int
import Foreign.Marshal.Alloc
import Foreign.Ptr
import Foreign.Storable
import GHC.Generics                           (Generic)
import System.Global6
import Network.Kademlia
import Network.Address                        (Address, WantIP (..), either4or6,
                                               fromSockAddr, ipFamily, setPort,
                                               sockAddrPort, testIdBit,
                                               toSockAddr, un4map)
import Network.Kademlia.Search          (Search (..))
import qualified Network.Kademlia.Routing          as R
import Network.QueryResponse
import Network.Socket
import System.Endian
import System.IO
import qualified Text.ParserCombinators.ReadP as RP
import Text.Printf
import Text.Read
import Control.TriadCommittee
import Network.BitTorrent.DHT.Token           as Token
import GHC.TypeLits

import Connection
import Crypto.Tox
import Data.Word64Map                        (fitsInInt)
import qualified Data.Word64Map              (empty)
import Network.Kademlia.Bootstrap            (forkPollForRefresh, bootstrap)
import Network.Tox.Crypto.Transport          (NetCrypto(..), CryptoMessage, HandshakeData(..), Handshake(..))
import Network.Tox.Crypto.Handlers
import qualified Network.Tox.DHT.Handlers    as DHT
import qualified Network.Tox.DHT.Transport   as DHT
import Network.Tox.NodeId
import qualified Network.Tox.Onion.Handlers  as Onion
import qualified Network.Tox.Onion.Transport as Onion
import Network.Tox.Transport
import OnionRouter
import Network.Tox.ContactInfo
import Text.XXD
import qualified Data.HashMap.Strict as HashMap
import Data.HashMap.Strict (HashMap)
import qualified Data.Map.Strict as Map
import Control.Concurrent (threadDelay)
import DPut


newCrypto :: IO TransportCrypto
newCrypto = do
    secret <- generateSecretKey
    alias <- generateSecretKey
    ralias <- generateSecretKey
    let pubkey    = toPublic secret
        aliaspub  = toPublic alias
        raliaspub = toPublic ralias
    ukeys <- atomically $ newTVar []
    (symkey, drg) <- do
        drg0 <- getSystemDRG
        return $ randomBytesGenerate 32 drg0 :: IO (ByteString, SystemDRG)
    noncevar <- atomically $ newTVar $ fst $ withDRG drg drgNew
    cookieKeys <- atomically $ newTVar []
    cache <- newSecretsCache
    hPutStrLn stderr $ "secret(tox) = " ++ DHT.showHex secret
    hPutStrLn stderr $ "public(tox) = " ++ DHT.showHex pubkey
    hPutStrLn stderr $ "symmetric(tox) = " ++ DHT.showHex symkey
    return TransportCrypto
        { transportSecret = secret
        , transportPublic = pubkey
        , onionAliasSecret = alias
        , onionAliasPublic = aliaspub
        , rendezvousSecret = ralias
        , rendezvousPublic = raliaspub
        , transportSymmetric = return $ SymmetricKey symkey
        , transportNewNonce = do
            drg1 <- readTVar noncevar
            let (nonce, drg2) = withDRG drg1 (Nonce24 <$> getRandomBytes 24)
            writeTVar noncevar drg2
            return nonce
        , userKeys = return []
        , pendingCookies = cookieKeys
        , secretsCache = cache
        }

updateIP :: TVar (R.BucketList NodeInfo) -> SockAddr -> STM ()
updateIP tblvar a = do
    bkts <- readTVar tblvar
    case nodeInfo (nodeId (R.thisNode bkts)) a of
        Right ni -> writeTVar tblvar (bkts { R.thisNode = ni })
        Left _   -> return ()

genNonce24 :: DRG g =>
              TVar (g, pending) -> DHT.TransactionId -> IO DHT.TransactionId
genNonce24 var (DHT.TransactionId nonce8 _) = atomically $ do
    (g,pending) <- readTVar var
    let (bs, g') = randomBytesGenerate 24 g
    writeTVar var (g',pending)
    return $ DHT.TransactionId nonce8 (Nonce24 bs)


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

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

w64Key :: DHT.TransactionId -> Word64
w64Key (DHT.TransactionId (Nonce8 w) _) = w

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

-- | Return my own address.
myAddr :: TVar (R.BucketList NodeInfo) -- ^ IPv4 buckets
       -> TVar (R.BucketList NodeInfo) -- ^ IPv6 buckets
       -> Maybe NodeInfo             -- ^ Interested remote address
       -> IO NodeInfo
myAddr routing4 routing6 maddr = atomically $ do
    let var = case flip DHT.prefer4or6 Nothing <$> maddr of
                Just Want_IP6 -> routing4
                _             -> routing6
    a <- readTVar var
    return $ R.thisNode a

newClient :: (DRG g, Show addr, Show meth) =>
              g -> Transport String addr x
                -> (Client String meth DHT.TransactionId addr x -> x -> MessageClass String meth DHT.TransactionId addr x)
                -> (Maybe addr -> IO addr)
                -> (Client String meth DHT.TransactionId addr x -> meth -> Maybe (MethodHandler String DHT.TransactionId addr x))
                -> (forall d. TransactionMethods d DHT.TransactionId addr x -> TransactionMethods d DHT.TransactionId addr x)
                -> (Client String meth DHT.TransactionId addr x -> Transport String addr x -> Transport String addr x)
                -> IO (Client String meth DHT.TransactionId addr x)
newClient drg net classify selfAddr handlers modifytbl modifynet = do
    -- If we have 8-byte keys for IntMap, then use it for transaction lookups.
    -- Otherwise, use ordinary Map.  The details of which will be hidden by an
    -- existential closure (see mkclient below).
    --
    tblvar <-
      if fitsInInt (Proxy :: Proxy Word64)
        then do
            let intmapT = transactionMethods (contramap intKey intMapMethods) gen
            intmap_var <- atomically $ newTVar (drg, mempty)
            return $ Right (intmapT,intmap_var)
         else do
            let word64mapT = transactionMethods (contramap w64Key w64MapMethods) gen
            map_var <- atomically $ newTVar (drg, Data.Word64Map.empty)
            return $ Left (word64mapT,map_var)
    let dispatch tbl var handlers client = DispatchMethods
                { classifyInbound = classify client
                , lookupHandler = handlers -- var
                , tableMethods = modifytbl tbl
                }
        eprinter = printErrors stderr
        mkclient (tbl,var) handlers =
            let client = Client
                    { clientNet           = addHandler (reportParseError eprinter) (handleMessage client) $ modifynet client net
                    , clientDispatcher    = dispatch tbl var (handlers client) client
                    , clientErrorReporter = eprinter { reportTimeout = reportTimeout ignoreErrors }
                    , clientPending       = var
                    , clientAddress       = selfAddr
                    , clientResponseId    = genNonce24 var
                    }
            in client
    return $ either mkclient mkclient tblvar handlers

data Tox = Tox
    { toxDHT            :: DHT.Client
    , toxOnion          :: Onion.Client RouteId
    , toxToRoute        :: Transport String Onion.AnnouncedRendezvous (PublicKey,Onion.OnionData)
    , toxCrypto         :: Transport String SockAddr NetCrypto
    , toxCryptoSessions :: NetCryptoSessions
    , toxCryptoKeys     :: TransportCrypto
    , toxRouting        :: DHT.Routing
    , toxTokens         :: TVar SessionTokens
    , toxAnnouncedKeys  :: TVar Onion.AnnouncedKeys
    , toxOnionRoutes    :: OnionRouter
    , toxContactInfo    :: ContactInfo
    , toxAnnounceToLan  :: IO ()
    }

-- | initiate a netcrypto session, blocking
netCrypto :: Tox -> SecretKey -> PublicKey{-UserKey -} -> IO [NetCryptoSession]
netCrypto tox myseckey theirpubkey = netCryptoWithBackoff 1000000 tox myseckey theirpubkey

-- | helper for 'netCrypto', initiate a netcrypto session, retry after specified millisecs
netCryptoWithBackoff :: Int -> Tox -> SecretKey -> PublicKey -> IO [NetCryptoSession]
netCryptoWithBackoff millisecs tox myseckey theirpubkey = do
    let mykeyAsId = key2id (toPublic myseckey)
    mbContactsVar <- fmap contacts . HashMap.lookup mykeyAsId <$> atomically (readTVar (accounts (toxContactInfo tox)))
    case mbContactsVar of
      Nothing -> do
        hPutStrLn stderr ("netCrypto: (" ++ show mykeyAsId ++") accounts lookup failed.")
        return []

      Just contactsVar -> do
        let theirkeyAsId = key2id theirpubkey
        mbContact <- HashMap.lookup theirkeyAsId <$> atomically (readTVar contactsVar)
        case mbContact of
         Nothing -> do
            hPutStrLn stderr ("netCrypto: (" ++ show mykeyAsId ++") friend not found (" ++ show theirkeyAsId ++ ").")
            return []
         Just contact@(Contact { contactKeyPacket     = mbKeyPkt
                               , contactLastSeenAddr  = Nothing
                               , contactFriendRequest = mbFR
                               , contactPolicy        = mbPolicy
                               }) -> do
            hPutStrLn stderr ("netCrypto: (" ++ show mykeyAsId ++") no SockAddr for friend (" ++ show theirkeyAsId ++ "). TODO: search their node?")
            return []
         Just contact@(Contact { contactKeyPacket     = Nothing
                               }) -> do
            hPutStrLn stderr ("netCrypto: (" ++ show mykeyAsId ++") no DHT-key for friend (" ++ show theirkeyAsId ++ "). TODO: what?")
            return []
         Just contact@(Contact { contactKeyPacket     = Just keyPkt
                               , contactLastSeenAddr  = Just saddr
                               , contactFriendRequest = mbFR
                               , contactPolicy        = mbPolicy
                               }) | theirDhtKey <- DHT.dhtpk keyPkt -> do
            -- Do we already have an active session with this user?
            sessionsMap <- atomically $ readTVar (netCryptoSessionsByKey (toxCryptoSessions tox) )
            let sessionUsesIdentity key session = key == ncMyPublicKey session
            case Map.lookup theirpubkey sessionsMap of
            -- if sessions found, is it using this private key?
              Just sessions | matchedSessions <- filter (sessionUsesIdentity (toPublic myseckey)) sessions
                            , not (null matchedSessions)
                            -> do
                hPutStrLn stderr ("netCrypto: Already have a session for " ++ show mykeyAsId ++ "<-->" ++ show theirkeyAsId)
                return matchedSessions
            -- if not, send handshake, this is separate session
              Nothing -> do
                -- if no session:
                -- Convert to NodeInfo, so we can send cookieRequest
                let crypto = toxCryptoKeys tox
                    client = toxDHT tox
                case nodeInfo (key2id theirDhtKey) saddr of
                   Left e   -> hPutStrLn stderr ("netCrypto: nodeInfo fail... " ++ e) >> return []
                   Right ni -> do
                    mbCookie <- DHT.cookieRequest crypto client (toPublic myseckey) ni
                    case mbCookie of
                      Nothing -> do
                        hPutStrLn stderr ("netCrypto: (" ++ show mykeyAsId ++") <--> (" ++ show theirkeyAsId ++ ").")
                        hPutStrLn stderr ("netCrypto: CookieRequest failed. TODO: dhtpkNodes thingy")
                        return []
                      Just cookie -> do
                        hPutStrLn stderr "Have cookie, creating handshake packet..."
                        let hp = HParam { hpOtherCookie = cookie
                                        , hpMySecretKey = myseckey
                                        , hpCookieRemotePubkey = theirpubkey
                                        , hpCookieRemoteDhtkey = theirDhtKey
                                        , hpTheirBaseNonce = error "netCrypto: Unreachable! hpTheirBaseNonce"
                                        , hpTheirSessionKeyPublic = error "netCrypto: Unreachable! hpTheirSessionKeyPublic"
                                        }
                        myhandshake <- do
                            n24' <- atomically $ transportNewNonce crypto
                            dput XNetCrypto ("Handshake Nonce24: " <> show n24')
                            newBaseNonce <- atomically $ transportNewNonce crypto
                            mbMyhandshakeData <- newHandShakeData crypto newBaseNonce hp saddr
                            forM mbMyhandshakeData $ \hsdata -> do
                                state <- lookupSharedSecret crypto myseckey theirpubkey n24'
                                return Handshake { handshakeCookie = cookie
                                                 , handshakeNonce = n24'
                                                 , handshakeData = encrypt state $ encodePlain hsdata
                                                 }
                        case myhandshake of
                            Nothing -> hPutStrLn stderr "netCrypto: failed to create HandshakeData." >> return []
                            Just handshake -> do
                                sendMessage (toxCrypto tox) saddr (NetHandshake handshake)
                                let secnum :: Double
                                    secnum = fromIntegral millisecs / 1000000
                                    delay = (millisecs * 5 `div` 4)
                                if secnum < 20000000
                                  then do
                                    hPutStrLn stderr $ "sent handshake, now delaying " ++ show (secnum * 1.25) ++ "  second(s).."
                                    threadDelay delay
                                    netCryptoWithBackoff delay tox myseckey theirpubkey -- hopefully it will find an active session this time.
                                  else do
                                    hPutStrLn stderr "Unable to establish session..."
                                    return []

getContactInfo :: Tox -> IO DHT.DHTPublicKey
getContactInfo Tox{toxCryptoKeys,toxRouting} = join $ atomically $ do
    r4 <- readTVar $ DHT.routing4 toxRouting
    r6 <- readTVar $ DHT.routing6 toxRouting
    nonce <- transportNewNonce toxCryptoKeys
    let self = nodeId n4
        n4 = R.thisNode r4
        n6 = R.thisNode r6
        n4s = R.kclosest DHT.toxSpace 4 self r4
        n6s = R.kclosest DHT.toxSpace 4 self r6
        ns = filter (DHT.isGlobal . nodeIP) [n4,n6]
              ++ concat (zipWith (\a b -> [a,b]) n4s n6s)
    return $ do
        timestamp <- round . (* 1000000) <$> getPOSIXTime
        return DHT.DHTPublicKey
            { dhtpkNonce = timestamp
            , dhtpk      = id2key self
            , dhtpkNodes = DHT.SendNodes $ take 4 ns
            }

isLocalHost :: SockAddr -> Bool
isLocalHost (SockAddrInet _ host32) = (fromBE32 host32 == 0x7f000001)
isLocalHost _                       = False

addVerbosity :: Transport err SockAddr ByteString -> Transport err SockAddr ByteString
addVerbosity tr =
    tr { awaitMessage = \kont -> awaitMessage tr $ \m -> do
            forM_ m $ mapM_ $ \(msg,addr) -> do
                when (not (B.null msg || elem (B.head msg) [0,1,2,4,0x81,0x82,0x8c,0x8d])) $ do
                    mapM_ (\x -> hPutStrLn stderr ( (show addr) ++ " --> " ++ x))
                          $ xxd 0 msg
            kont m
       , sendMessage = \addr msg -> do
            when (not (B.null msg || elem (B.head msg) [0,1,2,4,0x81,0x8c,0x8d])) $ do
                mapM_ (\x -> hPutStrLn stderr ( (show addr) ++ " <-- " ++ x))
                      $ xxd 0 msg
            sendMessage tr addr msg
       }

newKeysDatabase :: IO (TVar Onion.AnnouncedKeys)
newKeysDatabase =
    atomically $ newTVar $ Onion.AnnouncedKeys PSQ.empty MinMaxPSQ.empty


getOnionAlias :: TransportCrypto -> STM NodeInfo -> Maybe (Onion.OnionDestination r) -> IO (Onion.OnionDestination r)
getOnionAlias crypto dhtself remoteNode = atomically $ do
    ni <- dhtself
    let alias = case remoteNode of
            Just (Onion.OnionDestination (Onion.AnnouncingAlias _ uk) _ _)
              -> ni { nodeId = key2id uk }
            _ -> ni { nodeId = key2id (onionAliasPublic crypto) }
    return $ Onion.OnionDestination Onion.SearchingAlias alias Nothing


newTox :: TVar Onion.AnnouncedKeys   -- ^ Store of announced keys we are a rendezvous for.
       -> SockAddr                   -- ^ Bind-address to listen on.
       -> Maybe NetCryptoSessions    -- ^ State of all one-on-one Tox links.
       -> Maybe SecretKey            -- ^ Optional DHT secret key to use.
       -> IO Tox
newTox keydb addr mbSessionsState suppliedDHTKey = do
    (udp,sock) <- {- addVerbosity <$> -} udpTransport' addr
    (crypto0,sessionsState0) <- case mbSessionsState of
                                 Nothing -> do
                                    crypto <- newCrypto
                                    sessionsState <- newSessionsState crypto defaultUnRecHook defaultCryptoDataHooks
                                    return (crypto,sessionsState)
                                 Just s -> return (transportCrypto s, s)

    roster <- newContactInfo
    let crypto = fromMaybe crypto0 $do
            k <- suppliedDHTKey
            return crypto0
                { transportSecret = k
                , transportPublic = toPublic k
                , userKeys = myKeyPairs roster
                }
    forM_ suppliedDHTKey $ \k -> do
        maybe (hPutStrLn stderr "failed to encode suppliedDHTKey")
              (C8.hPutStrLn stderr . C8.append "Using suppliedDHTKey: ")
              $ encodeSecret k

    drg <- drgNew
    let lookupClose _ = return Nothing

    mkrouting <- DHT.newRouting addr crypto updateIP updateIP
    let ignoreErrors _ = return () -- Set this to (hPutStrLn stderr) to debug onion route building.
    orouter <- newOnionRouter ignoreErrors
    (dhtcrypt,onioncrypt,dtacrypt,cryptonet) <- toxTransport crypto orouter lookupClose udp
    let sessionsState = sessionsState0 { sessionTransport = cryptonet
                                       , transportCrypto = crypto }
    let dhtnet0 = layerTransportM (DHT.decrypt crypto) (DHT.encrypt crypto) dhtcrypt
        tbl4 = DHT.routing4 $ mkrouting (error "missing client")
        tbl6 = DHT.routing6 $ mkrouting (error "missing client")
    dhtclient <- newClient drg dhtnet0 DHT.classify (myAddr tbl4 tbl6) (DHT.handlers crypto . mkrouting) id
                    $ \client net -> onInbound (DHT.updateRouting client (mkrouting client) orouter) net

    orouter <- forkRouteBuilder orouter $ \nid ni -> fmap (\(_,ns,_)->ns) <$> DHT.getNodes dhtclient nid ni

    toks <- do
        nil <- nullSessionTokens
        atomically $ newTVar nil { maxInterval = 20 } -- 20 second timeout on announce ping-ids.
    oniondrg <- drgNew
    let onionnet = layerTransportM (Onion.decrypt crypto) (Onion.encrypt crypto) onioncrypt
    onionclient <- newClient oniondrg onionnet (const Onion.classify)
                    (getOnionAlias crypto $ R.thisNode <$> readTVar (DHT.routing4 $ mkrouting dhtclient))
                    (const $ Onion.handlers onionnet (mkrouting dhtclient) toks keydb)
                    (hookQueries orouter DHT.transactionKey)
                    (const id)

    return Tox
        { toxDHT            = dhtclient
        , toxOnion          = onionclient
        , toxToRoute        = onInbound (updateContactInfo roster) dtacrypt
        , toxCrypto         = addHandler (hPutStrLn stderr) (cryptoNetHandler sessionsState) cryptonet
        , toxCryptoSessions = sessionsState
        , toxCryptoKeys     = crypto
        , toxRouting        = mkrouting dhtclient
        , toxTokens         = toks
        , toxAnnouncedKeys  = keydb
        , toxOnionRoutes    = orouter
        , toxContactInfo    = roster
        , toxAnnounceToLan  = announceToLan sock (key2id $ transportPublic crypto)
        }

onionTimeout :: Tox -> DHT.TransactionId -> Onion.OnionDestination RouteId -> STM (Onion.OnionDestination RouteId, Int)
onionTimeout Tox { toxOnionRoutes = or } (DHT.TransactionId n8 _) od = lookupTimeout or n8 od

forkTox :: Tox -> IO (IO (), [NodeInfo] -> [NodeInfo] -> IO (), [NodeInfo] -> [NodeInfo] -> IO ())
forkTox tox = do
    _ <- forkListener "toxCrypto" (toxCrypto tox)
    _ <- forkListener "toxToRoute" (toxToRoute tox)
    _ <- forkListener "toxOnion" (clientNet $ toxOnion tox)
    quit <- forkListener "toxDHT" (clientNet $ toxDHT tox)
    forkPollForRefresh (DHT.refresher4 $ toxRouting tox)
    forkPollForRefresh (DHT.refresher6 $ toxRouting tox)
    return ( quit
           , bootstrap (DHT.refresher4 $ toxRouting tox)
           , bootstrap (DHT.refresher6 $ toxRouting tox)
           )

-- TODO: Don't export this.
announceToLan :: Socket -> NodeId -> IO ()
announceToLan sock nid = do
    (broadcast_info:_) <- getAddrInfo (Just defaultHints { addrFlags = [AI_NUMERICHOST], addrSocketType = Datagram })
                                      (Just "192.168.1.255") -- TODO: Detect broadcast address.
                                      (Just "33445")
    let broadcast = addrAddress broadcast_info
        bs = S.runPut $ DHT.putMessage (DHT.DHTLanDiscovery nid)
    saferSendTo sock bs broadcast