{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveFunctor, DeriveTraversable,DeriveDataTypeable #-}
{-# LANGUAGE GADTs #-}
module ToxMessage where

import Data.ByteString (ByteString)
import qualified Crypto.MAC.Poly1305 as Poly1305 (Auth(..))
import qualified Crypto.PubKey.Curve25519 as Curve25519
import Data.ByteArray as BA (ByteArrayAccess, Bytes)
import qualified Data.ByteArray as BA
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as C8
import qualified Data.ByteString.Base16 as Base16
import Data.Bits
import Data.Hashable
import Data.Bits.ByteString ()
import Data.Word
import Data.Data
import Data.Ord
import Data.Serialize

newtype Auth = Auth Poly1305.Auth
 deriving (Eq, ByteArrayAccess)

instance Ord Auth where
    compare (Auth a) (Auth b) = comparing (BA.convert :: Poly1305.Auth -> Bytes) a b

instance Data Auth where
    gfoldl k z x = z x

    -- Well, this is a little wonky... XXX
    gunfold k z c = k (z (Auth . Poly1305.Auth . (BA.convert :: ByteString -> Bytes)))

    toConstr _ = con_Auth

    dataTypeOf _ = mkDataType "ToxMessage" [con_Auth]

con_Auth = mkConstr (dataTypeOf (Auth (error "con_Auth"))) "Auth" [] Prefix

instance Serialize Auth where
    get = Auth . Poly1305.Auth . BA.convert <$> getBytes 16
    put (Auth (Poly1305.Auth bs)) = putByteString $ BA.convert bs
--
-- | An 'Aliased' 'PubKey' is one that is not the DHT key and so should not go
-- into the kademlia routing table buckets.
--
-- Note: This includes the long-term tox-id key that humans use to friend each
-- other and is often refered to as your "real public key" by the Tox
-- documents.  For the purposes of the DHT, it is an alias.
newtype Aliased a = Aliased a
 deriving (Eq,Ord,Show,Data,Functor,Foldable,Traversable,Serialize)

newtype Nonce24 = Nonce24 ByteString
 deriving (Eq, Ord, ByteArrayAccess,Data)

quoted :: ShowS -> ShowS
quoted shows s = '"':shows ('"':s)

bin2hex :: ByteArrayAccess bs => bs -> String
bin2hex = C8.unpack . Base16.encode . BA.convert

instance Show Nonce24 where
    showsPrec d nonce = quoted (mappend $ bin2hex nonce)

instance Serialize Nonce24 where
    get = Nonce24 <$> getBytes 24
    put (Nonce24 bs) = putByteString bs

newtype Nonce8 = Nonce8 Word64
 deriving (Eq, Ord,Data)

-- TODO: This should probably be represented by Curve25519.PublicKey, but
-- ByteString has more instances...
newtype PubKey = PubKey ByteString
 deriving (Eq,Ord,Data, ByteArrayAccess, Bits, Hashable)

instance Serialize PubKey where
    get = PubKey <$> getBytes 32
    put (PubKey bs) = putByteString bs

instance Show PubKey where
    show (PubKey bs) = C8.unpack $ Base16.encode bs

instance FiniteBits PubKey where
    finiteBitSize _ = 256

instance Read PubKey where
    readsPrec _ str
        | (bs, xs) <- Base16.decode $ C8.pack str
        , B.length bs == 32
                    = [ (PubKey bs, drop 64 str) ]
        | otherwise = []



-- | A chunk of data encrypted with public-key cryptography.
data ImplicitAssymetric = ImplicitAssymetric
    { assymetricAuth  :: Auth
    , assymetricBytes :: ByteString
    }
 deriving (Eq, Ord,Data)

getRemaining = remaining >>= getBytes

instance Serialize ImplicitAssymetric where
    get = ImplicitAssymetric <$> get <*> getRemaining
    put (ImplicitAssymetric auth bs) = put auth >> putByteString bs

-- | Like ImplicitAssymetric, but includes the nonce used to encrypt.
data UnclaimedAssymetric = UnclaimedAssymetric
    { assymetricNonce :: Nonce24
    , assymetricData  :: !ImplicitAssymetric
    }
 deriving (Eq, Ord, Data)

instance Serialize UnclaimedAssymetric where
    get = UnclaimedAssymetric <$> get <*> get
    put (UnclaimedAssymetric nonce dta) = put nonce >> put dta

-- | Like UnclaimedAssymetric, but includes the public key of the sender.
data Assymetric = Assymetric
    { senderKey :: PubKey
    , sent      :: !UnclaimedAssymetric
    }
 deriving (Eq, Ord,Data)

-- get requires isolate.
instance Serialize Assymetric where
    get = Assymetric <$> get <*> get
    put (Assymetric key dta) = put key >> put dta

newtype Cookie = Cookie UnclaimedAssymetric
 deriving (Eq, Ord,Data)

data Symmetric = Symmetric
    { symmetricNonce :: Nonce24
    , symmetricAuth  :: Auth
    , symmetricBytes :: ByteString
    }
 deriving (Eq, Ord,Data)

data Packet where
    Ping          :: Assymetric -> Packet -- 0x00 -- Assymetric query
    Pong          :: Assymetric -> Packet -- 0x01 -- Assymetric response

    GetNodes      :: Assymetric -> Packet -- 0x02 -- Assymetric query
    SendNodes     :: Assymetric -> Packet -- 0x04 -- Assymetric response

    CookieRequest :: Assymetric -> Packet -- 0x18
    CookieResponse :: UnclaimedAssymetric -> Packet -- 0x19

    OnionRequest0 :: Assymetric -> Packet -- 0x80


    CryptoHandshake :: Cookie -> UnclaimedAssymetric -> Packet -- 0x1a

    CryptoData :: Word16 -> ImplicitAssymetric -> Packet -- 0x1b

    DHTRequest :: PubKey -> Assymetric -> Packet -- 0x20 -- Sometimes Assymetric query

    DataToRoute :: PubKey -> Aliased Assymetric -> Packet
    DataToRouteResponse :: Aliased Assymetric -> Packet

    LanDiscovery :: PubKey -> Packet  -- 0x21

    OnionRequest1 :: Aliased Assymetric -> Symmetric -> Packet -- 0x81
    OnionRequest2 :: Aliased Assymetric -> Symmetric -> Packet -- 0x82

    OnionRequest3 :: ByteString -> Symmetric -> Packet -- 0x82

    Announce :: Aliased Assymetric -> Packet --0x83
    AnnounceResponse :: Nonce8 -> UnclaimedAssymetric -> Packet -- 0x84

    OnionResponse3 :: Symmetric -> ByteString -> Packet -- 0x8c
    OnionResponse2 :: Symmetric -> ByteString -> Packet -- 0x8d
    OnionResponse1 :: Symmetric -> ByteString -> Packet -- 0x8e


    GetBootstrapInfo :: ByteString -> Packet -- 0xf0 + 77 bytes                         -- ByteString query
    BootstrapInfo :: Word32 -> ByteString -> Packet -- 0xf0 + version + (256 byte motd) -- ByteSTring response

 deriving (Eq, Ord,Data)

newtype PacketKind = PacketKind Word8
 deriving (Eq, Ord, Serialize)

pktKind :: Packet -> PacketKind

-- These are (Assymetric -> Assymetric) queries.
pktKind Ping             {} = PacketKind 0x00
pktKind Pong             {} = PacketKind 0x01
pktKind GetNodes         {} = PacketKind 0x02
pktKind SendNodes        {} = PacketKind 0x04


-- This is a (Assymetric -> UnclaimedAssymetric) query
pktKind CookieRequest    {} = PacketKind 0x18
pktKind CookieResponse   {} = PacketKind 0x19

-- Query (Assymetric -> (Nonce8,UnclaimedAssymetric))
pktKind Announce         {} = PacketKind 0x83
pktKind AnnounceResponse {} = PacketKind 0x84

-- Query (Assymetric -> ByteString)
pktKind OnionRequest0    {} = PacketKind 0x80


-- This is a (ByteString -> ByteString) query
pktKind GetBootstrapInfo {} = PacketKind 0xf0
pktKind BootstrapInfo    {} = PacketKind 0xf0


-- Trigering event. No direct response. (PubKey -> ())
pktKind LanDiscovery     {} = PacketKind 0x21

-- Two cases:
--  Half-established:    (Cookie,UnclaimedAssymetric) -> (Cookie,UnclaimedAssymetric)
--  Session established: (Cookie,UnclaimedAssymetric) -> (Word16,ImplicitAssymetric)
pktKind CryptoHandshake  {} = PacketKind 0x1a

-- Session data, no direct response.
--   (reponse to CryptoHandshake, or other data)
pktKind CryptoData       {} = PacketKind 0x1b

-- Two cases:
--    ( (PubKey, Assymetric) -> response )
--    ( (PubKey, Assymetric) -> ()       )
pktKind DHTRequest       {} = PacketKind 0x20


-- Query ( (PubKey,Aliased Assymetric) -> Aliased Assymetric)
pktKind DataToRoute      {}    = PacketKind 0x85
pktKind DataToRouteResponse {} = PacketKind 0x86

-- 3 Queries ( (Aliased Assymetric, Symmetric )
--            -> ( Symmetric, ByteString )     )
pktKind OnionRequest1    {} = PacketKind 0x81
pktKind OnionResponse1   {} = PacketKind 0x8e

pktKind OnionRequest2    {} = PacketKind 0x82
pktKind OnionResponse2   {} = PacketKind 0x8d

pktKind OnionRequest3    {} = PacketKind 0x82
pktKind OnionResponse3   {} = PacketKind 0x8c

data PacketClass =
    AssymetricClass (Assymetric -> Packet) (Packet -> Assymetric)
    | Unclassified

pktClass :: PacketKind -> PacketClass
pktClass (PacketKind 0) = AssymetricClass Ping (\(Ping a) -> a)
pktClass (PacketKind 1) = AssymetricClass Pong (\(Pong a) -> a)
pktClass (PacketKind 2) = AssymetricClass GetNodes (\(GetNodes a) -> a)
pktClass (PacketKind 4) = AssymetricClass SendNodes (\(SendNodes a) -> a)

pktClass (PacketKind 0x18) = AssymetricClass CookieRequest (\(CookieRequest a) -> a)
pktClass (PacketKind 0x80) = AssymetricClass OnionRequest0 (\(OnionRequest0 a) -> a)
pktClass (PacketKind 0x86) = AssymetricClass (DataToRouteResponse . Aliased) (\(DataToRouteResponse (Aliased a)) -> a)

    -- (indexConstr (dataTypeOf (error "dataTypeOf Packet" :: Packet)) 0) -- Ping

instance Serialize Packet where
    get = getPacket
    put = putPacket

getPacket = do
    typ <- get
    case pktClass typ of
        AssymetricClass toPacket fromPacket -> toPacket <$> get

putPacket p = do
    put $ pktKind p
    case pktClass (pktKind p) of
        AssymetricClass toPacket fromPacket -> put $ fromPacket p

{-
data Packet' where
    :: Assymetric -> Packet
    :: UnclaimedAssymetric -> Packet
    :: Word16 -> ImplicitAssymetric -> Packet
    :: PubKey -> Assymetric -> Packet
    :: PubKey -> Packet
    :: Aliased Assymetric -> Symmetric -> Packet
    :: ByteString -> Symmetric -> Packet
    :: Aliased Assymetric -> Packet
    :: Symmetric -> ByteString -> Packet
    :: ByteString -> Packet
    :: Word32 -> ByteString -> Packet
-}