{-# LANGUAGE PatternSynonyms            #-}
module DHTHandlers where

import DHTTransport
import Network.QueryResponse as QR hiding (Client)
import qualified Network.QueryResponse as QR (Client)
import ToxCrypto
import ToxMessage  as Tox (PacketKind(..), pattern PingType, pattern GetNodesType, pattern DHTRequestType)

import qualified Data.Wrapper.PSQInt as Int
import Kademlia
import Network.Address               (WantIP (..), ipFamily, testIdBit)
import qualified Network.DHT.Routing as R
import TriadCommittee

import Control.Monad
import Control.Concurrent.STM
import Data.Time.Clock.POSIX  (POSIXTime, getPOSIXTime)
import Network.Socket
import Data.IP
import Data.Maybe
import Data.Bits

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

classify :: DHTMessage ((,) Nonce8) -> MessageClass String PacketKind TransactionId
classify msg = mapMessage (\nonce24 (nonce8,_) -> go msg (TransactionId nonce8 nonce24)) msg
 where
    go (DHTPing          {}) = IsQuery PingType
    go (DHTGetNodes      {}) = IsQuery GetNodesType
    go (DHTPong          {}) = IsResponse
    go (DHTSendNodes     {}) = IsResponse
    go (DHTCookieRequest {}) = IsQuery (PacketKind 0x18)
    go (DHTCookie        {}) = IsResponse
    go (DHTDHTRequest    {}) = IsQuery DHTRequestType

data Routing = Routing
    { tentativeId :: NodeInfo
    , sched4      :: !( TVar (Int.PSQ POSIXTime) )
    , routing4    :: !( TVar (R.BucketList NodeInfo) )
    , committee4  :: TriadCommittee NodeId SockAddr
    , sched6      :: !( TVar (Int.PSQ POSIXTime) )
    , routing6    :: !( TVar (R.BucketList NodeInfo) )
    , committee6  :: TriadCommittee NodeId SockAddr
    }


-- TODO: This should cover more cases
isLocal (IPv6 ip6) = (ip6 == toEnum 0)
isLocal (IPv4 ip4) = (ip4 == toEnum 0)

isGlobal = not . isLocal

prefer4or6 :: NodeInfo -> Maybe WantIP -> WantIP
prefer4or6 addr iptyp = fromMaybe (ipFamily $ nodeIP addr) iptyp

toxSpace :: R.KademliaSpace NodeId NodeInfo
toxSpace = R.KademliaSpace
    { R.kademliaLocation = nodeId
    , R.kademliaTestBit  = testIdBit
    , R.kademliaXor      = xor
    }


pingH :: NodeInfo -> Ping -> IO Pong
pingH _ Ping = return Pong

getNodesH :: Routing -> NodeInfo -> GetNodes -> IO SendNodes
getNodesH routing addr (GetNodes nid) = do
    let preferred = prefer4or6 addr Nothing

    (append4,append6) <- atomically $ do
        ni4 <- R.thisNode <$> readTVar (routing4 routing)
        ni6 <- R.thisNode <$> readTVar (routing6 routing)
        return $ case ipFamily (nodeIP addr) of
            Want_IP4 | isGlobal (nodeIP ni6) -> (id, (++ [ni6]))
            Want_IP6 | isGlobal (nodeIP ni4) -> ((++ [ni4]), id)
            _                                -> (id, id)
    ks  <- go append4 $ routing4 routing
    ks6 <- go append6 $ routing6 routing
    let (ns1,ns2) = case preferred of Want_IP6 -> (ks6,ks)
                                      Want_IP4 -> (ks,ks6)
    return $ SendNodes
           $ if null ns2 then ns1
                         else take 4 (take 3 ns1 ++ ns2)
 where
    go f var = f . R.kclosest toxSpace k nid <$> atomically (readTVar var)

    k = 4

type Message = DHTMessage ((,) Nonce8)

type Client = QR.Client String PacketKind TransactionId NodeInfo Message


wrapAssym :: TransactionId -> NodeInfo -> NodeInfo -> dta -> Assym (Nonce8,dta)
wrapAssym (TransactionId n8 n24) src dst dta = Assym
    { senderKey = let NodeId pubkey = nodeId src in pubkey
    , assymNonce = n24
    , assymData = (n8, dta)
    }

serializer :: PacketKind
                -> (Assym (Nonce8,ping) -> Message)
                -> (Message -> Maybe (Assym (Nonce8,pong)))
                -> MethodSerializer TransactionId NodeInfo Message PacketKind ping (Maybe pong)
serializer pktkind mkping mkpong = MethodSerializer
        { methodTimeout  = 5
        , method         = pktkind
        -- wrapQuery :: tid -> addr -> addr -> qry -> x
        , wrapQuery      = \tid src dst ping -> mkping $ wrapAssym tid src dst ping
        -- unwrapResponse :: x -> b
        , unwrapResponse = fmap (snd . assymData) . mkpong
        }


unpong :: Message -> Maybe (Assym (Nonce8,Pong))
unpong (DHTPong assym) = Just assym
unpong _               = Nothing

ping :: Client -> NodeInfo -> IO Bool
ping client addr = do
    reply <- QR.sendQuery client (serializer PingType DHTPing unpong) Ping addr
    maybe (return False) (\Pong -> return True) $ join reply

unsendNodes :: Message -> Maybe (Assym (Nonce8,SendNodes))
unsendNodes (DHTSendNodes assym) = Just assym
unsendNodes _                    = Nothing

unwrapNodes :: SendNodes -> ( [NodeInfo], [NodeInfo], () )
unwrapNodes (SendNodes ns) = (ns,ns,())

getNodes :: Client -> NodeId -> NodeInfo -> IO (Maybe ([NodeInfo],[NodeInfo],()))
getNodes client nid addr = do
    reply <- QR.sendQuery client (serializer GetNodesType DHTGetNodes unsendNodes) (GetNodes nid) addr
    return $ fmap unwrapNodes $ join reply

updateRouting :: Client -> Routing -> NodeInfo -> Message -> IO ()
updateRouting client routing naddr msg = do
    -- hPutStrLn stderr $ "updateRouting "++show typ
    -- TODO: check msg type
    case prefer4or6 naddr Nothing of
        Want_IP4 -> updateTable client naddr (routing4 routing) (committee4 routing) (sched4 routing)
        Want_IP6 -> updateTable client naddr (routing6 routing) (committee6 routing) (sched6 routing)

updateTable :: Client -> NodeInfo -> TVar (R.BucketList NodeInfo) -> TriadCommittee NodeId SockAddr -> TVar (Int.PSQ POSIXTime) -> IO ()
updateTable client naddr tbl committee sched = do
        self <- atomically $ R.thisNode <$> readTVar tbl
        when (nodeIP self /= nodeIP naddr) $ do
            -- TODO: IP address vote?
            insertNode (toxKademlia client committee tbl sched) naddr

toxKademlia :: Client -> TriadCommittee NodeId SockAddr -> TVar (R.BucketList NodeInfo) -> TVar (Int.PSQ POSIXTime) -> Kademlia NodeId NodeInfo
toxKademlia client committee var sched
    = Kademlia quietInsertions
               toxSpace
               (vanillaIO var $ ping client)
                   { tblTransition = \tr -> do
                        io1 <- transitionCommittee committee tr
                        io2 <- touchBucket toxSpace (15*60) var sched tr
                        return $ do
                            io1 >> io2
                            {-
                            hPutStrLn stderr $ unwords
                                [ show (transitionedTo tr)
                                , show (transitioningNode tr)
                                ]
                            -}
                            return ()
                   }

transitionCommittee :: TriadCommittee NodeId SockAddr -> RoutingTransition NodeInfo -> STM (IO ())
transitionCommittee committee (RoutingTransition ni Stranger) = do
    delVote committee (nodeId ni)
    return $ do
        -- hPutStrLn stderr $ "delVote "++show (nodeId ni)
        return ()
transitionCommittee committee _ = return $ return ()