{-# LANGUAGE RecursiveDo #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE FlexibleContexts #-}
module Network.Tox.TCP
    ( module Network.Tox.TCP
    , NodeInfo(..)
    ) where

import Debug.Trace
import Control.Arrow
import Control.Concurrent
import Control.Concurrent.STM
import Control.Exception
import Control.Monad
import Crypto.Random
import Data.Aeson (ToJSON(..),FromJSON(..))
import qualified Data.Aeson as JSON
import Data.Functor.Contravariant
import Data.Functor.Identity
import Data.Hashable
import qualified Data.HashMap.Strict as HashMap
import Data.IP
import Data.Maybe
import Data.Monoid
import Data.Serialize
import Data.Word
import qualified Data.Vector as Vector
import Network.Socket (SockAddr(..))
import qualified Text.ParserCombinators.ReadP as RP
import System.IO.Error
import System.Timeout

import ControlMaybe
import Crypto.Tox
import Data.ByteString (hPut,hGet,ByteString,length)
import Data.TableMethods
import Data.Tox.Relay
import qualified Data.Word64Map
import DebugTag
import DPut
import Network.Address (setPort,PortNumber,localhost4,fromSockAddr)
import Network.Kademlia.Routing
import Network.Kademlia.Search hiding (sendQuery)
import Network.QueryResponse
import Network.QueryResponse.TCP
import Network.Tox.DHT.Handlers (toxSpace)
import Network.Tox.Onion.Transport hiding (encrypt,decrypt)
import Network.Tox.Onion.Handlers (unwrapAnnounceResponse)
import qualified Network.Tox.NodeId as UDP
import Text.XXD
import Data.Proxy

withSize :: Sized x => (Size x -> m (p x)) -> m (p x)
withSize f = case size of len -> f len


type NodeId = UDP.NodeId

-- example:
-- KEyW2Bm.S-DpIGp72380BAfgintUWX1KX.6ZU.4m5Ex@80.99.99.99:33400{tcp:443}
instance Show NodeInfo where
    show (NodeInfo udp port) = show udp ++ "{tcp:"++show port++"}"

nodeId :: NodeInfo -> NodeId
nodeId ni = UDP.nodeId $ udpNodeInfo ni

nodeAddr :: NodeInfo -> SockAddr
nodeAddr ni = setPort (tcpPort ni) $ UDP.nodeAddr $ udpNodeInfo ni

nodeIP :: NodeInfo -> IP
nodeIP ni = UDP.nodeIP $ udpNodeInfo ni

tcpStream :: (Show y, Show x, Serialize y, Sized y, Serialize x, Sized x) =>
                   TransportCrypto -> StreamHandshake NodeInfo x y
tcpStream crypto = StreamHandshake
    { streamHello = \addr h -> do
        (skey, hello) <- atomically $ do
            n24 <- transportNewNonce crypto
            skey <- transportNewKey crypto
            base24 <- transportNewNonce crypto
            return $ (,) skey $ Hello $ Asymm
                { senderKey  = transportPublic crypto
                , asymmNonce = n24
                , asymmData  = pure HelloData
                                { sessionPublicKey = toPublic $ skey
                                , sessionBaseNonce = base24
                                }
                }
        noncef <- lookupNonceFunction crypto (transportSecret crypto) (UDP.id2key $ nodeId addr)
        dput XTCP $ "TCP:" ++ show addr ++ " <-- " ++ show hello
        hPut h $ encode $ encryptPayload (noncef $ helloNonce hello) hello
        (welcomeE, wbs) <- do
            let sz0 = size
                sz = constSize sz0
            bs <- hGet h sz
            return ( fmap (`asProxyTypeOf` sz0) $ decode bs, bs )
        let mwelcome = welcomeE >>= \w -> decryptPayload (noncef $ welcomeNonce w) w
            nil = SessionProtocol
                { streamGoodbye = return ()
                , streamDecode = return Nothing
                , streamEncode = \y -> dput XTCP $ "TCP nil <-- " ++ show y
                }
        either (\e -> do
                        dput XTCP $ "welcome: " ++ show (Data.ByteString.length wbs) ++ " bytes."
                        forM_ (xxd2 0 wbs) $ dput XTCP
                        dput XTCP $ "TCP(fail welcome): " ++ e
                        return nil
                ) id $ mwelcome <&> \welcome -> do
        dput XTCP $ "TCP:" ++ show addr ++ " --> " ++ show welcome
        noncef' <- lookupNonceFunction crypto skey (sessionPublicKey $ runIdentity $ welcomeData welcome)
        nsend <- newMVar (sessionBaseNonce $ runIdentity $ helloData hello)
        nread <- newMVar (sessionBaseNonce $ runIdentity $ welcomeData welcome)
        let them = sessionPublicKey $ runIdentity $ welcomeData welcome
        hvar <- newMVar h
        return SessionProtocol
            { streamGoodbye = do
                dput XTCP $ "Closing " ++ show addr
                return () -- No goodbye packet?  Seems rude.
            , streamDecode =
                let go h = decode <$> hGet h 2 >>= \case
                        Left e    -> do
                            dput XTCP $ "TCP: (" ++ show addr ++ ") Failed to get length: " ++ e
                            return Nothing
                        Right len -> do
                            decode <$> hGet h (fromIntegral (len :: Word16)) >>= \case
                                Left e  -> do
                                    dput XTCP $ "TCP: Failed to decode packet."
                                    return Nothing
                                Right x -> do
                                    dput XTCP $ "TCP:"++ show addr ++ " --> packet!"
                                    m24 <- timeout 1000000 (takeMVar nread)
                                    fmap join $ forM m24 $ \n24 -> do
                                        let r = decrypt (noncef' n24) x >>= decodePlain
                                        putMVar nread (incrementNonce24 n24)
                                        either (dput XTCP . ("TCP decryption: " ++))
                                               (\x' -> do
                                                    dput XTCP $ "TCP:" ++ show addr ++ " --> " ++ show x'
                                                    return ())
                                               r
                                        return $ either (const Nothing) Just r
                 in bracket (takeMVar hvar) (putMVar hvar)
                     $ \h -> go h `catchIOError` \e -> do
                                        dput XTCP $ "TCP exception: " ++ show e
                                        return Nothing
            , streamEncode = \y -> do
                -- dput XTCP $ "TCP(acquire nonce):" ++ show addr ++ " <-- " ++ show y
                n24 <- takeMVar nsend
                -- dput XTCP $ "TCP(got nonce):" ++ show addr ++ " <-- " ++ show y
                let bs = encode $ encrypt (noncef' n24) $ encodePlain y
                ($ h) -- bracket (takeMVar hvar) (putMVar hvar)
                    $ \h -> hPut h (encode (fromIntegral $ Data.ByteString.length bs :: Word16) <> bs)
                                `catchIOError` \e -> dput XTCP $ "TCP write exception: " ++ show e
                -- dput XTCP $ "TCP(incrementing nonce): " ++ show addr ++ " <-- " ++ show y
                putMVar nsend (incrementNonce24 n24)
                dput XTCP $ "TCP: " ++ show addr ++ " <-- " ++ show y
            }
    , streamAddr = nodeAddr
    }

toxTCP :: TransportCrypto -> IO ( TCPCache (SessionProtocol RelayPacket RelayPacket)
                                , TransportA err NodeInfo RelayPacket (Bool,RelayPacket) )
toxTCP crypto = tcpTransport 30 (tcpStream crypto)

tcpSpace :: KademliaSpace NodeId NodeInfo
tcpSpace = contramap udpNodeInfo toxSpace

{-
nodeSearch :: TCPClient err () Nonce8 -> Search NodeId (IP, PortNumber) () NodeInfo NodeInfo
nodeSearch tcp = Search
    { searchSpace       = tcpSpace
    , searchNodeAddress = nodeIP &&& tcpPort
    , searchQuery       = getNodes tcp
    }
-}

data TCPClient err tid = TCPClient
    { tcpCrypto     :: TransportCrypto
    , tcpClient     :: Client err PacketNumber tid NodeInfo (Bool,RelayPacket)
    , tcpGetGateway :: UDP.NodeInfo -> STM (Maybe NodeInfo)
    }

{-
getTCPNodes :: TCPClient err () Nonce8 -> NodeId -> NodeInfo -> IO (Maybe ([NodeInfo],[NodeInfo],Maybe ()))
getTCPNodes tcp seeking dst = do
    r <- getUDPNodes' tcp seeking (udpNodeInfo dst)
    let tcps (ns,_,mb) = (ns',ns',mb)
         where ns' = do
                    n <- ns
                    [ NodeInfo n (fromIntegral 443) , NodeInfo n (fromIntegral 80) , NodeInfo n (UDP.nodePort n) ]
    fmap join $ forM r $ \(ns,gw) -> do
        let ts = tcps ns
        {-
        if nodeId gw == nodeId dst
            then return $ Just ts
            else do
                forkIO $ void $ tcpPing (tcpClient tcp) dst
                return $ Just ts
        -}
        forM_ ((\(xs,_,_) -> xs) ts) (forkIO . void . tcpPing (tcpClient tcp))
        return $ Just ts
-}

getUDPNodes :: TCPClient err Nonce8 -> NodeId -> UDP.NodeInfo -> IO (Maybe ([UDP.NodeInfo], [UDP.NodeInfo], Maybe ()))
getUDPNodes tcp seeking dst = fmap fst <$> getUDPNodes' tcp seeking dst

getUDPNodes' :: TCPClient err Nonce8 -> NodeId -> UDP.NodeInfo -> IO (Maybe (([UDP.NodeInfo], [UDP.NodeInfo], Maybe ()), NodeInfo))
getUDPNodes' tcp seeking dst0 = do
    mgateway <- atomically $ tcpGetGateway tcp dst0
    fmap join $ forM mgateway $ \gateway -> do
        (b,c,n24) <- atomically $ do
            b <- transportNewKey (tcpCrypto tcp)
            c <- transportNewKey (tcpCrypto tcp)
            n24 <- transportNewNonce (tcpCrypto tcp)
            return (b,c,n24)
        let (dst,gateway') = if UDP.nodeId dst0 == nodeId gateway
                                then ( dst0 { UDP.nodeIP = fromJust $ Network.Address.fromSockAddr localhost4 }
                                     , gateway { udpNodeInfo = (udpNodeInfo gateway)
                                                    { UDP.nodeIP = fromJust $ Network.Address.fromSockAddr localhost4 }})
                                else (dst0,gateway)
        wrap2 <- lookupNonceFunction (tcpCrypto tcp) b (UDP.id2key $ UDP.nodeId dst)
        wrap1 <- lookupNonceFunction (tcpCrypto tcp) c (UDP.id2key $ nodeId gateway)
        wrap0 <- lookupNonceFunction (tcpCrypto tcp) (transportSecret $ tcpCrypto tcp) (UDP.id2key $ UDP.nodeId dst)
        let meth :: MethodSerializer
                        Nonce8
                        a -- NodeInfo
                        (Bool, RelayPacket)
                        PacketNumber
                        AnnounceRequest
                        (Either String AnnounceResponse)
            meth = MethodSerializer
                    { methodTimeout = \tid addr -> return (addr,12000000) -- 12 second timeout
                    , method = OnionPacketID -- meth
                    , wrapQuery = \n8 src gateway x -> (,) True $
                        OnionPacket n24 $ Addressed (UDP.nodeAddr dst)
                            $ wrapOnionPure b (wrap2 n24) (UDP.nodeAddr $ udpNodeInfo gateway')
                                $ wrapOnionPure c (wrap1 n24) (UDP.nodeAddr dst)
                                    $ NotForwarded $ encryptPayload (wrap0 n24)
                                       $ OnionAnnounce Asymm
                                            { senderKey  = transportPublic (tcpCrypto tcp)
                                            , asymmNonce = n24
                                            , asymmData  = pure (x,n8)
                                            }
                    , unwrapResponse = \case
                        (_,OnionPacketResponse (OnionAnnounceResponse _ n24' r))
                            -> decrypt (wrap0 n24') r >>= decodePlain
                        x   -> Left $ "getUDPNodes: unwrapResponse fail " ++ show x
                    }
        r <- sendQuery (tcpClient tcp) meth (AnnounceRequest zeros32 seeking UDP.zeroID) gateway
        forM r $ \response -> do
            let (ns,_,mb) = either (const ([],[],Nothing)) (unwrapAnnounceResponse Nothing dst) $ response
            return ( (ns,ns, const () <$> mb), gateway )


handleOOB :: PublicKey -> ByteString -> NodeInfo -> NodeInfo -> IO (Maybe (x -> x))
handleOOB k bs src dst = do
    dput XMisc $ "TODO: handleOOB " ++ show src
    return Nothing

handle2route :: OnionMessage Encrypted -> NodeInfo -> NodeInfo -> IO (Maybe (x -> x))
handle2route o src dst = do
    dput XMisc $ "TODO: handle2route " ++ show src
    return Nothing

tcpPing :: Show addr => Client err PacketNumber Nonce8 addr (Bool,RelayPacket) -> addr -> IO (Maybe ())
tcpPing client dst = do
    dput XTCP $ "tcpPing " ++ show dst
    sendQuery client meth () dst
 where meth = MethodSerializer
                { wrapQuery      = \n8 src dst () -> (True,RelayPing n8)
                , unwrapResponse = \_ -> ()
                , methodTimeout  = \n8 dst -> return (dst,5000000)
                , method         = PingPacket
                }

type RelayClient = Client String PacketNumber Nonce8 NodeInfo (Bool,RelayPacket)

-- | Create a new TCP relay client.  Because polymorphic existential record
-- updates are currently hard with GHC, this function accepts parameters for
-- generalizing the table-entry type for pending transactions.  Safe trivial
-- defaults are 'id' and 'tryPutMVar'.  The resulting customized table state
-- will be returned to the caller along with the new client.
newClient :: TransportCrypto
               -> ((Maybe (Bool,RelayPacket) -> IO ()) -> a) -- ^ store mvar for query
               -> (a -> RelayPacket -> IO void)  -- ^ load mvar for query
               -> IO ( ( TVar (ChaChaDRG, Data.Word64Map.Word64Map a)
                       , TCPCache (SessionProtocol RelayPacket RelayPacket) )
                     , Client String PacketNumber Nonce8 NodeInfo (Bool,RelayPacket))
newClient crypto store load = do
    (tcpcache,net) <- toxTCP crypto
    drg <- drgNew
    map_var <- atomically $ newTVar (drg, Data.Word64Map.empty)
    return $ (,) (map_var,tcpcache) Client
        { clientNet = {- XXX: Client type forces this pointless layering. -} layerTransport ((Right .) . (,) . (,) False) (,) net
        , clientDispatcher = DispatchMethods
            { classifyInbound = (. snd) $ \case
                RelayPing n           -> IsQuery PingPacket n
                RelayPong n           -> IsResponse n
                OnionPacketResponse (OnionAnnounceResponse n8 n24 ciphered) -> IsResponse n8
                OnionPacketResponse o@(OnionToRouteResponse _) -> IsUnsolicited $ handle2route o
                OOBRecv k bs          -> IsUnsolicited $ handleOOB k bs
                wut                   -> IsUnknown (show wut)
            , lookupHandler   = \case
                PingPacket -> trace ("tcp-received-ping") $ Just MethodHandler
                    { methodParse     = \case (_,RelayPing n8) -> Right ()
                                              _                -> trace ("tcp-non-ping") $ Left "TCP: Non-ping?"
                    , methodSerialize = \n8 src dst () -> {- trace ("tcp-made-pong-"++show n8) -} (False, RelayPong n8)
                    , methodAction    = \src () -> dput XTCP $ "TCP pinged by "++show src
                    }
                w -> trace ("tcp-lookupHandler: "++show w) $ Just NoReply
                            { methodParse = \x -> Left "tcp-lookuphandler?" -- :: x -> Either err a
                            , noreplyAction = \addr a -> dput XTCP $ "tcp-lookupHandler: "++show w
                            }
            , tableMethods = transactionMethods' store (\x -> mapM_ (load x . snd)) (contramap (\(Nonce8 w64) -> w64) w64MapMethods)
                                $ first (either error Nonce8 . decode) . randomBytesGenerate 8
            }
        , clientErrorReporter = logErrors
        , clientPending = map_var
        , clientAddress = \_ -> return $ NodeInfo
            { udpNodeInfo = either error id $ UDP.nodeInfo (UDP.key2id $ transportPublic crypto) (SockAddrInet 0 0)
            , tcpPort = 0
            }
        , clientResponseId = return
        }