{-# LANGUAGE RecursiveDo #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE FlexibleContexts #-}
module Network.Tox.TCP where

import Control.Arrow
import Control.Concurrent
import Control.Concurrent.STM
import Crypto.Random
import Data.Functor.Contravariant
import Data.Functor.Identity
import Data.IP
import Data.Serialize
import Network.Socket (SockAddr(..))

import Crypto.Tox
import Data.ByteString (hPut,hGet,ByteString)
import Data.Tox.Relay
import qualified Data.Word64Map
import DebugTag
import DPut
import Network.Address (setPort,PortNumber)
import Network.Kademlia.Routing
import Network.Kademlia.Search
import Network.QueryResponse
import Network.QueryResponse.TCP
import Network.Tox.DHT.Handlers (toxSpace)
import Network.Tox.Onion.Transport (OnionMessage(..),OnionResponse(..),N1)
import qualified Network.Tox.NodeId as UDP


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

data NodeInfo = NodeInfo
    { udpNodeInfo :: UDP.NodeInfo
    , tcpPort     :: PortNumber
    }

type NodeId = UDP.NodeId

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 :: (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)
        hPut h $ encode $ encryptPayload (noncef $ helloNonce hello) hello
        welcomeE <- withSize $ fmap decode . hGet h . constSize
        let Right welcome = welcomeE >>= \w -> decryptPayload (noncef $ welcomeNonce w) w
        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
        return SessionProtocol
            { streamGoodbye = return () -- No goodbye packet?  Seems rude.
            , streamDecode = do
                decode <$> hGet h 2 >>= \case
                    Left _    -> return Nothing
                    Right len -> do
                        decode <$> hGet h len >>= \case
                            Left _  -> return Nothing
                            Right x -> do
                                n24 <- takeMVar nread
                                let r = decrypt (noncef' n24) x >>= decodePlain
                                putMVar nread (incrementNonce24 n24)
                                return $ either (const Nothing) Just r
            , streamEncode = \y -> do
                n24 <- takeMVar nsend
                hPut h $ encode $ encrypt (noncef' n24) $ encodePlain y
                putMVar nsend (incrementNonce24 n24)
            }
    , streamAddr = nodeAddr
    }

toxTCP :: TransportCrypto -> IO (Transport err NodeInfo RelayPacket)
toxTCP crypto = tcpTransport 30 (tcpStream crypto)

tcpSpace :: KademliaSpace NodeId NodeInfo
tcpSpace = contramap udpNodeInfo toxSpace

nodeSearch :: Client err meth tid NodeInfo RelayPacket -> Search NodeId (IP, PortNumber) () NodeInfo NodeInfo
nodeSearch client = Search
    { searchSpace       = tcpSpace
    , searchNodeAddress = nodeIP &&& tcpPort
    , searchQuery       = getNodes client
    }

getNodes :: Client err meth tid NodeInfo RelayPacket -> NodeId -> NodeInfo -> IO (Maybe ([NodeInfo], [NodeInfo], Maybe ()))
getNodes client seeking dst = do
    return Nothing -- TODO

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

handle2route :: OnionResponse N1 -> NodeInfo -> NodeInfo -> IO (Maybe (RelayPacket -> RelayPacket))
handle2route o src dst = do
    dput XMisc $ "TODO: handle2route " ++ show src
    return Nothing


tcpClient :: TransportCrypto -> IO (Client String () Nonce8 NodeInfo RelayPacket)
tcpClient crypto = do
    net <- toxTCP crypto
    drg <- drgNew
    map_var <- atomically $ newTVar (drg, Data.Word64Map.empty)
    return Client
        { clientNet = net
        , clientDispatcher = DispatchMethods
            { classifyInbound = \case
                RelayPing n           -> IsQuery () n
                RelayPong n           -> IsResponse n
                OnionPacketResponse (OnionResponse _ (OnionAnnounceResponse n8 n24 ciphered)) -> IsResponse n8
                OnionPacketResponse o@(OnionResponse _ (OnionToRouteResponse _)) -> IsUnsolicited $ handle2route o
                OOBRecv k bs          -> IsUnsolicited $ handleOOB k bs
            , lookupHandler   = \() -> Just MethodHandler
                    { methodParse     = \(RelayPing n8) -> Right ()
                    , methodSerialize = \n8 src dst () -> RelayPong n8
                    , methodAction    = \src () -> return ()
                    }
            , tableMethods = transactionMethods (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
        , clientEnterQuery = \_ -> return ()
        , clientLeaveQuery = \_ _ -> return ()
        }