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

import Control.Concurrent
import Control.Concurrent.STM
import Data.Functor.Identity
import Data.Serialize
import System.IO (Handle)

import Crypto.Tox
import Data.ByteString (hPut,hGet,ByteString)
import Data.Tox.Relay
import Network.Address (setPort,PortNumber,SockAddr)
import Network.QueryResponse
import Network.QueryResponse.TCP
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

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

nodeAddr :: NodeInfo -> SockAddr
nodeAddr ni = setPort (tcpPort ni) $ UDP.nodeAddr $ 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 = \h -> return () -- No goodbye packet?  Seems rude.
            , streamDecode = \h -> 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
                let bs = encode $ encrypt (noncef' n24) $ encodePlain y
                putMVar nsend (incrementNonce24 n24)
                return bs -- XXX: Should we wait until this bytestring is sent before putting the nonce back in the MVar?
            }
    , streamAddr = nodeAddr
    }

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