{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE NoFieldSelectors #-}

-- |
-- Module      : Haskoin.Network.Message
-- Copyright   : No rights reserved
-- License     : MIT
-- Maintainer  : jprupp@protonmail.ch
-- Stability   : experimental
-- Portability : POSIX
--
-- Peer-to-peer network message serialization.
module Haskoin.Network.Message
  ( -- * Network Message
    Message (..),
    MessageHeader (..),
    msgType,
    putMessage,
    getMessage,
  )
where

import Control.DeepSeq
import Control.Monad (unless)
import Data.Binary (Binary (..))
import Data.ByteString (ByteString)
import Data.ByteString qualified as B
import Data.Bytes.Get
import Data.Bytes.Put
import Data.Bytes.Serial
import Data.Serialize (Serialize (..))
import Data.Word (Word32)
import GHC.Generics (Generic)
import Haskoin.Block.Common
import Haskoin.Block.Merkle
import Haskoin.Crypto.Hash
import Haskoin.Network.Bloom
import Haskoin.Network.Common
import Haskoin.Network.Data
import Haskoin.Transaction.Common

-- | Data type representing the header of a 'Message'. All messages sent between
-- nodes contain a message header.
data MessageHeader = MessageHeader
  { -- | magic bytes identify network
    magic :: !Word32,
    -- | message type
    cmd :: !MessageCommand,
    -- | length of payload
    size :: !Word32,
    -- | checksum of payload
    checksum :: !CheckSum32
  }
  deriving (Eq, Show, Generic, NFData)

instance Serial MessageHeader where
  deserialize =
    MessageHeader
      <$> getWord32be
      <*> deserialize
      <*> getWord32le
      <*> deserialize

  serialize (MessageHeader m c l chk) = do
    putWord32be m
    serialize c
    putWord32le l
    serialize chk

instance Binary MessageHeader where
  put = serialize
  get = deserialize

instance Serialize MessageHeader where
  put = serialize
  get = deserialize

-- | The 'Message' type is used to identify all the valid messages that can be
-- sent between bitcoin peers. Only values of type 'Message' will be accepted
-- by other bitcoin peers as bitcoin protocol messages need to be correctly
-- serialized with message headers. Serializing a 'Message' value will
-- include the 'MessageHeader' with the correct checksum value automatically.
-- No need to add the 'MessageHeader' separately.
data Message
  = MVersion !Version
  | MVerAck
  | MAddr !Addr
  | MInv !Inv
  | MGetData !GetData
  | MNotFound !NotFound
  | MGetBlocks !GetBlocks
  | MGetHeaders !GetHeaders
  | MTx !Tx
  | MBlock !Block
  | MMerkleBlock !MerkleBlock
  | MHeaders !Headers
  | MGetAddr
  | MFilterLoad !FilterLoad
  | MFilterAdd !FilterAdd
  | MFilterClear
  | MPing !Ping
  | MPong !Pong
  | MAlert !Alert
  | MMempool
  | MReject !Reject
  | MSendHeaders
  | MOther !ByteString !ByteString
  deriving (Eq, Show, Generic, NFData)

-- | Get 'MessageCommand' assocated with a message.
msgType :: Message -> MessageCommand
msgType (MVersion _) = MCVersion
msgType MVerAck = MCVerAck
msgType (MAddr _) = MCAddr
msgType (MInv _) = MCInv
msgType (MGetData _) = MCGetData
msgType (MNotFound _) = MCNotFound
msgType (MGetBlocks _) = MCGetBlocks
msgType (MGetHeaders _) = MCGetHeaders
msgType (MTx _) = MCTx
msgType (MBlock _) = MCBlock
msgType (MMerkleBlock _) = MCMerkleBlock
msgType (MHeaders _) = MCHeaders
msgType (MFilterLoad _) = MCFilterLoad
msgType (MFilterAdd _) = MCFilterAdd
msgType MFilterClear = MCFilterClear
msgType (MPing _) = MCPing
msgType (MPong _) = MCPong
msgType (MAlert _) = MCAlert
msgType MMempool = MCMempool
msgType (MReject _) = MCReject
msgType MSendHeaders = MCSendHeaders
msgType MGetAddr = MCGetAddr
msgType (MOther c _) = MCOther c

-- | Deserializer for network messages.
getMessage :: (MonadGet m) => Network -> m Message
getMessage net = do
  (MessageHeader mgc cmd len chk) <- deserialize
  bs <- lookAhead $ getByteString $ fromIntegral len
  unless
    (mgc == net.magic)
    (fail $ "get: Invalid network magic bytes: " ++ show mgc)
  unless
    (checkSum32 bs == chk)
    (fail $ "get: Invalid message checksum: " ++ show chk)
  if len > 0
    then do
      bs <- ensure (fromIntegral len)
      let f = case cmd of
            MCVersion -> MVersion <$> deserialize
            MCAddr -> MAddr <$> deserialize
            MCInv -> MInv <$> deserialize
            MCGetData -> MGetData <$> deserialize
            MCNotFound -> MNotFound <$> deserialize
            MCGetBlocks -> MGetBlocks <$> deserialize
            MCGetHeaders -> MGetHeaders <$> deserialize
            MCTx -> MTx <$> deserialize
            MCBlock -> MBlock <$> deserialize
            MCMerkleBlock -> MMerkleBlock <$> deserialize
            MCHeaders -> MHeaders <$> deserialize
            MCFilterLoad -> MFilterLoad <$> deserialize
            MCFilterAdd -> MFilterAdd <$> deserialize
            MCPing -> MPing <$> deserialize
            MCPong -> MPong <$> deserialize
            MCAlert -> MAlert <$> deserialize
            MCReject -> MReject <$> deserialize
            MCOther c -> MOther c <$> getByteString (fromIntegral len)
            _ ->
              fail $
                "get: command "
                  ++ show cmd
                  ++ " should not carry a payload"
      either fail return (runGetS f bs)
    else case cmd of
      MCGetAddr -> return MGetAddr
      MCVerAck -> return MVerAck
      MCFilterClear -> return MFilterClear
      MCMempool -> return MMempool
      MCSendHeaders -> return MSendHeaders
      MCOther c -> return (MOther c B.empty)
      _ ->
        fail $
          "get: command "
            ++ show cmd
            ++ " is expected to carry a payload"

-- | Serializer for network messages.
putMessage :: (MonadPut m) => Network -> Message -> m ()
putMessage net msg = do
  let (cmd, payload) =
        case msg of
          MVersion m -> (MCVersion, runPutS $ serialize m)
          MVerAck -> (MCVerAck, B.empty)
          MAddr m -> (MCAddr, runPutS $ serialize m)
          MInv m -> (MCInv, runPutS $ serialize m)
          MGetData m -> (MCGetData, runPutS $ serialize m)
          MNotFound m -> (MCNotFound, runPutS $ serialize m)
          MGetBlocks m -> (MCGetBlocks, runPutS $ serialize m)
          MGetHeaders m -> (MCGetHeaders, runPutS $ serialize m)
          MTx m -> (MCTx, runPutS $ serialize m)
          MBlock m -> (MCBlock, runPutS $ serialize m)
          MMerkleBlock m -> (MCMerkleBlock, runPutS $ serialize m)
          MHeaders m -> (MCHeaders, runPutS $ serialize m)
          MGetAddr -> (MCGetAddr, B.empty)
          MFilterLoad m -> (MCFilterLoad, runPutS $ serialize m)
          MFilterAdd m -> (MCFilterAdd, runPutS $ serialize m)
          MFilterClear -> (MCFilterClear, B.empty)
          MPing m -> (MCPing, runPutS $ serialize m)
          MPong m -> (MCPong, runPutS $ serialize m)
          MAlert m -> (MCAlert, runPutS $ serialize m)
          MMempool -> (MCMempool, B.empty)
          MReject m -> (MCReject, runPutS $ serialize m)
          MSendHeaders -> (MCSendHeaders, B.empty)
          MOther c p -> (MCOther c, p)
      chk = checkSum32 payload
      len = fromIntegral $ B.length payload
      header = MessageHeader net.magic cmd len chk
  serialize header
  putByteString payload