{-# LANGUAGE ConstraintKinds            #-}
{-# LANGUAGE DeriveDataTypeable         #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
module Data.Tox.Relay where

import Data.ByteString as B
import Data.Data
import Data.Functor.Contravariant
import Data.Monoid
import Data.Serialize
import Data.Word
import qualified Rank2

import Crypto.Tox
import Network.Tox.Onion.Transport

newtype ConId = ConId Word8
    deriving (Eq,Show,Ord,Data,Serialize)

badcon :: ConId
badcon = ConId 0

-- Maps to a range -120 .. 119
c2key :: ConId -> Maybe Int
c2key (ConId x) | x < 16    = Nothing
                | otherwise = Just $ case divMod (x - 15) 2 of
                                (q,0) -> negate $ fromIntegral q
                                (q,1) -> fromIntegral q

-- Maps to range 16 .. 255
-- negatives become odds
key2c :: Int -> ConId
key2c y = ConId $ if y < 0 then 15 + fromIntegral (negate y * 2)
                           else 16 + fromIntegral (y * 2)

data RelayPacket
    = RoutingRequest PublicKey
    | RoutingResponse ConId PublicKey -- 0 for refusal, 16-255 for success.
    | ConnectNotification ConId
    | DisconnectNotification ConId
    | RelayPing Word64
    | RelayPong Word64
    | OOBSend PublicKey ByteString
    | OOBRecv PublicKey ByteString
    | OnionPacket (OnionRequest N0)
    | OnionPacketResponse (OnionResponse N1)
    -- 0x0A through 0x0F reserved for future use.
    | RelayData ConId ByteString -- Word8 is a connection id.  Encoded as number 16 to 255.
 deriving (Eq,Ord,Show,Data)

packetNumber :: RelayPacket -> Word8
packetNumber (RelayData (ConId conid) _) = conid -- 0 to 15 not allowed.
packetNumber rp                          = fromIntegral $ pred $ constrIndex $ toConstr rp

instance Sized RelayPacket where
    size = mappend (ConstSize 1) $ VarSize $ \x -> case x of
        RoutingRequest k             -> 32
        RoutingResponse rpid k       -> 33
        ConnectNotification conid    -> 1
        DisconnectNotification conid -> 1
        RelayPing pingid             -> 8
        RelayPong pingid             -> 8
        OOBSend k bs                 -> 32 + B.length bs
        OOBRecv k bs                 -> 32 + B.length bs
        OnionPacket query            -> case contramap (`asTypeOf` query) size of
                                            ConstSize n -> n
                                            VarSize f   -> f query
        OnionPacketResponse answer   -> case contramap (`asTypeOf` answer) size of
                                            ConstSize n -> n
                                            VarSize f   -> f answer
        RelayData _ bs               -> B.length bs

instance Serialize RelayPacket where

    get = do
        pktid <- getWord8
        case pktid of
            0 -> RoutingRequest <$> getPublicKey
            1 -> RoutingResponse <$> get <*> getPublicKey
            2 -> ConnectNotification <$> get
            3 -> DisconnectNotification <$> get
            4 -> RelayPing <$> getWord64be
            5 -> RelayPong <$> getWord64be
            6 -> OOBSend <$> getPublicKey <*> (remaining >>= getBytes)
            7 -> OOBRecv <$> getPublicKey <*> (remaining >>= getBytes)
            8 -> OnionPacket <$> get
            9 -> OnionPacketResponse <$> get
            conid -> RelayData (ConId conid) <$> (remaining >>= getBytes)

    put rp = do
        putWord8 $ packetNumber rp
        case rp of
            RoutingRequest k             -> putPublicKey k
            RoutingResponse rpid k       -> put rpid >> putPublicKey k
            ConnectNotification conid    -> put conid
            DisconnectNotification conid -> put conid
            RelayPing pingid             -> putWord64be pingid
            RelayPong pingid             -> putWord64be pingid
            OOBSend k bs                 -> putPublicKey k >> putByteString bs
            OOBRecv k bs                 -> putPublicKey k >> putByteString bs
            OnionPacket query            -> put query
            OnionPacketResponse answer   -> put answer
            RelayData _ bs               -> putByteString bs

-- | Initial client-to-server handshake message.
newtype Hello (f :: * -> *) = Hello (Asymm (f HelloData))

helloFrom :: Hello f -> PublicKey
helloFrom (Hello x) = senderKey x

helloNonce :: Hello f -> Nonce24
helloNonce (Hello x) = asymmNonce x

helloData :: Hello f -> f HelloData
helloData (Hello x) = asymmData x

instance Rank2.Functor Hello where
    f <$> Hello (Asymm k n dta) = Hello $ Asymm k n (f dta)

instance Payload Serialize Hello where
    mapPayload _ f (Hello (Asymm k n dta)) = Hello $ Asymm k n (f dta)

instance Rank2.Foldable Hello where
    foldMap f (Hello (Asymm k n dta)) = f dta

instance Rank2.Traversable Hello where
    traverse f (Hello (Asymm k n dta)) = Hello . Asymm k n <$> f dta

instance Sized (Hello Encrypted) where
    size = ConstSize 56 <> contramap helloData size

instance Serialize (Hello Encrypted) where
    get = Hello <$> getAsymm
    put (Hello asym) = putAsymm asym

data HelloData = HelloData
    { sessionPublicKey :: PublicKey
    , sessionBaseNonce :: Nonce24
    }

instance Sized HelloData where size = ConstSize 56

instance Serialize HelloData where
    get = HelloData <$> getPublicKey <*> get
    put (HelloData k n) = putPublicKey k >> put n

-- | Handshake server-to-client response packet.
data Welcome (f :: * -> *) = Welcome
    { welcomeNonce :: Nonce24
    , welcomeData  :: f HelloData
    }

instance Rank2.Functor Welcome where
    f <$> Welcome n dta = Welcome n (f dta)

instance Payload Serialize Welcome where
    mapPayload _ f (Welcome n dta) = Welcome n (f dta)

instance Rank2.Foldable Welcome where
    foldMap f (Welcome _ dta) = f dta

instance Rank2.Traversable Welcome where
    traverse f (Welcome n dta) = Welcome n <$> f dta

instance Sized (Welcome Encrypted) where
    size = ConstSize 24 <> contramap welcomeData size

instance Serialize (Welcome Encrypted) where
    get = Welcome <$> get <*> get
    put (Welcome n dta) = put n >> put dta