{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE NoFieldSelectors #-}

-- |
-- Module      : Haskoin.Script.SigHash
-- Copyright   : No rights reserved
-- License     : MIT
-- Maintainer  : jprupp@protonmail.ch
-- Stability   : experimental
-- Portability : POSIX
--
-- Transaction signatures and related functions.
module Haskoin.Script.SigHash
  ( -- * Script Signatures
    SigHash (..),
    SigHashFlag (..),
    sigHashAll,
    sigHashNone,
    sigHashSingle,
    anyoneCanPay,
    hasForkIdFlag,
    setAnyoneCanPay,
    setForkIdFlag,
    isSigHashAll,
    isSigHashNone,
    isSigHashSingle,
    isSigHashUnknown,
    sigHashAddForkId,
    sigHashGetForkId,
    sigHashAddNetworkId,
    txSigHash,
    txSigHashForkId,
    TxSignature (..),
    decodeTxSig,
    encodeTxSig,
  )
where

import Control.DeepSeq
import Control.Monad
import Crypto.Secp256k1
import Data.Aeson
import Data.Bits
import Data.Bool (bool)
import Data.ByteString (ByteString)
import Data.ByteString qualified as B
import Data.Bytes.Get
import Data.Bytes.Put
import Data.Bytes.Serial
import Data.Hashable
import Data.Maybe
import Data.Scientific
import Data.Word
import GHC.Generics (Generic)
import Haskoin.Crypto.Hash
import Haskoin.Crypto.Signature
import Haskoin.Network.Common
import Haskoin.Network.Data
import Haskoin.Script.Common
import Haskoin.Transaction.Common
import Haskoin.Util

-- | Constant representing a SIGHASH flag that controls what is being signed.
data SigHashFlag
  = -- | sign all outputs
    SIGHASH_ALL
  | -- | sign no outputs
    SIGHASH_NONE
  | -- | sign the output index corresponding to the input
    SIGHASH_SINGLE
  | -- | replay protection for Bitcoin Cash transactions
    SIGHASH_FORKID
  | -- | new inputs can be added
    SIGHASH_ANYONECANPAY
  deriving (Eq, Ord, Show, Read, Generic)

instance NFData SigHashFlag

instance Hashable SigHashFlag

instance Enum SigHashFlag where
  fromEnum SIGHASH_ALL = 0x01
  fromEnum SIGHASH_NONE = 0x02
  fromEnum SIGHASH_SINGLE = 0x03
  fromEnum SIGHASH_FORKID = 0x40
  fromEnum SIGHASH_ANYONECANPAY = 0x80
  toEnum 0x01 = SIGHASH_ALL
  toEnum 0x02 = SIGHASH_NONE
  toEnum 0x03 = SIGHASH_SINGLE
  toEnum 0x40 = SIGHASH_FORKID
  toEnum 0x80 = SIGHASH_ANYONECANPAY
  toEnum _ = error "Not a valid sighash flag"

-- | Data type representing the different ways a transaction can be signed.
-- When producing a signature, a hash of the transaction is used as the message
-- to be signed. The 'SigHash' parameter controls which parts of the
-- transaction are used or ignored to produce the transaction hash. The idea is
-- that if some part of a transaction is not used to produce the transaction
-- hash, then you can change that part of the transaction after producing a
-- signature without invalidating that signature.
--
-- If the 'SIGHASH_ANYONECANPAY' flag is set (true), then only the current input
-- is signed. Otherwise, all of the inputs of a transaction are signed. The
-- default value for 'SIGHASH_ANYONECANPAY' is unset (false).
newtype SigHash
  = SigHash Word32
  deriving (Eq, Ord, Enum, Show, Read, Generic)
  deriving newtype (Bits, Integral, Num, Real, Hashable, NFData)

instance FromJSON SigHash where
  parseJSON =
    withScientific "sighash" $
      maybe mzero (return . SigHash) . toBoundedInteger

instance ToJSON SigHash where
  toJSON = Number . fromIntegral
  toEncoding (SigHash n) = toEncoding n

-- | SIGHASH_NONE as a byte.
sigHashNone :: SigHash
sigHashNone = fromIntegral $ fromEnum SIGHASH_NONE

-- | SIGHASH_ALL as a byte.
sigHashAll :: SigHash
sigHashAll = fromIntegral $ fromEnum SIGHASH_ALL

-- | SIGHASH_SINGLE as a byte.
sigHashSingle :: SigHash
sigHashSingle = fromIntegral $ fromEnum SIGHASH_SINGLE

-- | SIGHASH_FORKID as a byte.
sigHashForkId :: SigHash
sigHashForkId = fromIntegral $ fromEnum SIGHASH_FORKID

-- | SIGHASH_ANYONECANPAY as a byte.
sigHashAnyoneCanPay :: SigHash
sigHashAnyoneCanPay = fromIntegral $ fromEnum SIGHASH_ANYONECANPAY

-- | Set SIGHASH_FORKID flag.
setForkIdFlag :: SigHash -> SigHash
setForkIdFlag = (.|. sigHashForkId)

-- | Set SIGHASH_ANYONECANPAY flag.
setAnyoneCanPay :: SigHash -> SigHash
setAnyoneCanPay = (.|. sigHashAnyoneCanPay)

-- | Is the SIGHASH_FORKID flag set?
hasForkIdFlag :: SigHash -> Bool
hasForkIdFlag = (/= 0) . (.&. sigHashForkId)

-- | Is the SIGHASH_ANYONECANPAY flag set?
anyoneCanPay :: SigHash -> Bool
anyoneCanPay = (/= 0) . (.&. sigHashAnyoneCanPay)

-- | Returns 'True' if the 'SigHash' has the value 'SIGHASH_ALL'.
isSigHashAll :: SigHash -> Bool
isSigHashAll = (== sigHashAll) . (.&. 0x1f)

-- | Returns 'True' if the 'SigHash' has the value 'SIGHASH_NONE'.
isSigHashNone :: SigHash -> Bool
isSigHashNone = (== sigHashNone) . (.&. 0x1f)

-- | Returns 'True' if the 'SigHash' has the value 'SIGHASH_SINGLE'.
isSigHashSingle :: SigHash -> Bool
isSigHashSingle = (== sigHashSingle) . (.&. 0x1f)

-- | Returns 'True' if the 'SigHash' has the value 'SIGHASH_UNKNOWN'.
isSigHashUnknown :: SigHash -> Bool
isSigHashUnknown =
  (`notElem` [sigHashAll, sigHashNone, sigHashSingle]) . (.&. 0x1f)

-- | Add a fork id to a 'SigHash'.
sigHashAddForkId :: SigHash -> Word32 -> SigHash
sigHashAddForkId sh w = (fromIntegral w `shiftL` 8) .|. (sh .&. 0x000000ff)

-- | Add fork id of a particular network to a 'SigHash'.
sigHashAddNetworkId :: Network -> SigHash -> SigHash
sigHashAddNetworkId net =
  (`sigHashAddForkId` fromMaybe 0 net.sigHashForkId)

-- | Get fork id from 'SigHash'.
sigHashGetForkId :: SigHash -> Word32
sigHashGetForkId (SigHash n) = fromIntegral $ n `shiftR` 8

-- | Computes the hash that will be used for signing a transaction.
txSigHash ::
  Network ->
  -- | transaction to sign
  Tx ->
  -- | script from output being spent
  Script ->
  -- | value of output being spent
  Word64 ->
  -- | index of input being signed
  Int ->
  -- | what to sign
  SigHash ->
  -- | hash to be signed
  Hash256
txSigHash net tx out v i sh
  | hasForkIdFlag sh && isJust net.sigHashForkId =
      txSigHashForkId net tx out v i sh
  | otherwise = do
      let newIn = buildInputs tx.inputs fout i sh
      -- When SigSingle and input index > outputs, then sign integer 1
      fromMaybe one $ do
        newOut <- buildOutputs tx.outputs i sh
        let newTx = Tx tx.version newIn newOut [] tx.locktime
        return . doubleSHA256 . runPutS $ do
          serialize newTx
          putWord32le $ fromIntegral sh
  where
    fout = Script $ filter (/= OP_CODESEPARATOR) out.ops
    one = "0100000000000000000000000000000000000000000000000000000000000000"

-- | Build transaction inputs for computing sighashes.
buildInputs :: [TxIn] -> Script -> Int -> SigHash -> [TxIn]
buildInputs txins out i sh
  | anyoneCanPay sh = [serialOut (txins !! i)]
  | isSigHashAll sh || isSigHashUnknown sh = single
  | otherwise = zipWith noSeq single [0 ..]
  where
    serialOut TxIn {..} = TxIn {script = runPutS $ serialize out, ..}
    emptyIn TxIn {..} = TxIn {script = B.empty, ..}
    emptyIns = map emptyIn txins
    single = updateIndex i emptyIns serialOut
    noSeq TxIn {..} j = TxIn {sequence = if i == j then sequence else 0, ..}

-- | Build transaction outputs for computing sighashes.
buildOutputs :: [TxOut] -> Int -> SigHash -> Maybe [TxOut]
buildOutputs txos i sh
  | isSigHashAll sh || isSigHashUnknown sh = return txos
  | isSigHashNone sh = return []
  | i >= length txos = Nothing
  | otherwise = return $ buffer ++ [txos !! i]
  where
    buffer = replicate i $ TxOut maxBound B.empty

-- | Compute the hash that will be used for signing a transaction. This
-- function is used when the 'SIGHASH_FORKID' flag is set.
txSigHashForkId ::
  Network ->
  -- | transaction to sign
  Tx ->
  -- | script from output being spent
  Script ->
  -- | value of output being spent
  Word64 ->
  -- | index of input being signed
  Int ->
  -- | what to sign
  SigHash ->
  -- | hash to be signed
  Hash256
txSigHashForkId net tx out v i sh =
  doubleSHA256 . runPutS $ do
    putWord32le tx.version
    serialize hashPrevouts
    serialize hashSequence
    serialize (tx.inputs !! i).outpoint
    putScript out
    putWord64le v
    putWord32le (tx.inputs !! i).sequence
    serialize hashOutputs
    putWord32le tx.locktime
    putWord32le $ fromIntegral $ sigHashAddNetworkId net sh
  where
    hashPrevouts
      | not (anyoneCanPay sh) =
          doubleSHA256 . runPutS $ mapM_ (serialize . (.outpoint)) tx.inputs
      | otherwise = zeros
    hashSequence
      | not (anyoneCanPay sh || isSigHashSingle sh || isSigHashNone sh) =
          doubleSHA256 . runPutS $ mapM_ (putWord32le . (.sequence)) tx.inputs
      | otherwise = zeros
    hashOutputs
      | not (isSigHashSingle sh || isSigHashNone sh) =
          doubleSHA256 . runPutS $ mapM_ serialize tx.outputs
      | isSigHashSingle sh && i < length tx.outputs =
          doubleSHA256 . runPutS $ serialize $ tx.outputs !! i
      | otherwise = zeros
    putScript s = do
      let encodedScript = runPutS $ serialize s
      putVarInt $ B.length encodedScript
      putByteString encodedScript
    zeros :: Hash256
    zeros = "0000000000000000000000000000000000000000000000000000000000000000"

-- | Data type representing a signature together with a 'SigHash'. The 'SigHash'
-- is serialized as one byte at the end of an ECDSA 'Sig'. All signatures in
-- transaction inputs are of type 'TxSignature'.
data TxSignature
  = TxSignature
      { sig :: !Sig,
        hash :: !SigHash
      }
  | TxSignatureEmpty
  deriving (Eq, Show, Read, Generic, NFData)

instance Marshal (Network, Ctx) TxSignature where
  marshalPut (net, ctx) TxSignatureEmpty = return ()
  marshalPut (net, ctx) (TxSignature sig (SigHash n)) = do
    marshalPut ctx sig
    putWord8 (fromIntegral n)

  marshalGet (net, ctx) =
    bool decode empty =<< isEmpty
    where
      empty = return TxSignatureEmpty
      decode = do
        sig <- marshalGet ctx
        sh <- fromIntegral <$> getWord8
        when (isSigHashUnknown sh) $
          fail "Non-canonical signature: unknown hashtype byte"
        when (isNothing net.sigHashForkId && hasForkIdFlag sh) $
          fail "Non-canonical signature: invalid network for forkId"
        return $ TxSignature sig sh

instance MarshalJSON (Network, Ctx) TxSignature where
  marshalValue (net, ctx) = String . encodeHex . encodeTxSig net ctx
  marshalEncoding s = hexEncoding . runPutL . marshalPut s
  unmarshalValue (net, ctx) =
    withText "TxSignature" $ \t ->
      case decodeHex t of
        Nothing -> fail "Cannot decode hex signature"
        Just b -> case decodeTxSig net ctx b of
          Left e -> fail e
          Right s -> return s

encodeTxSig :: Network -> Ctx -> TxSignature -> ByteString
encodeTxSig net ctx = runPutS . marshalPut (net, ctx)

decodeTxSig :: Network -> Ctx -> ByteString -> Either String TxSignature
decodeTxSig net ctx =
  runGetS $ do
    sig <- marshalGet (net, ctx)
    e <- isEmpty
    unless e $
      fail "Non-canonical signature: multiple hashtype bytes"
    return sig