{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE NoFieldSelectors #-}

-- |
-- Module      : Haskoin.Transaction.Builder.Sign
-- Copyright   : No rights reserved
-- License     : MIT
-- Maintainer  : jprupp@protonmail.ch
-- Stability   : experimental
-- Portability : POSIX
--
-- Types and logic for signing transactions.
module Haskoin.Transaction.Builder.Sign
  ( SigInput (..),
    makeSignature,
    makeSigHash,
    signTx,
    findInputIndex,
    signInput,
    buildInput,
    sigKeys,
  )
where

import Control.DeepSeq
import Control.Monad
import Crypto.Secp256k1
import Data.Aeson
import Data.Aeson.Encoding
import Data.Bytes.Get
import Data.Bytes.Put
import Data.Bytes.Serial
import Data.Either
import Data.Hashable
import Data.List
import Data.Maybe
import Data.Word
import GHC.Generics
import Haskoin.Address
import Haskoin.Crypto.Hash
import Haskoin.Crypto.Keys.Common
import Haskoin.Crypto.Signature
import Haskoin.Network.Data
import Haskoin.Script
import Haskoin.Transaction.Common
import Haskoin.Transaction.Segwit
import Haskoin.Util

-- | Data type used to specify the signing parameters of a transaction input.
-- To sign an input, the previous output script, outpoint and sighash are
-- required. When signing a pay to script hash output, an additional redeem
-- script is required.
data SigInput = SigInput
  { -- | output script to spend
    -- ^ output script value
    script :: !ScriptOutput,
    -- | output script value
    -- ^ outpoint to spend
    value :: !Word64,
    -- | outpoint to spend
    -- ^ signature type
    outpoint :: !OutPoint,
    -- | signature type
    -- ^ redeem script
    sighash :: !SigHash,
    -- | redeem script
    redeem :: !(Maybe RedeemScript)
  }
  deriving (Show, Read, Eq, Generic, NFData)

instance MarshalJSON Ctx SigInput where
  marshalValue ctx (SigInput s v o h r) =
    object $
      [ "pkscript" .= marshalValue ctx s,
        "value" .= v,
        "outpoint" .= o,
        "sighash" .= h
      ]
        ++ [ "redeem" .= marshalValue ctx r
             | r <- maybeToList r
           ]

  marshalEncoding ctx (SigInput s v o h r) =
    pairs $
      mconcat
        [ "pkscript" `pair` marshalEncoding ctx s,
          "value" `pair` word64 v,
          "outpoint" `pair` toEncoding o,
          "sighash" `pair` toEncoding h,
          maybe mempty (pair "redeem" . marshalEncoding ctx) r
        ]

  unmarshalValue ctx =
    withObject "SigInput" $ \o ->
      SigInput
        <$> (unmarshalValue ctx =<< o .: "pkscript")
        <*> o .: "value"
        <*> o .: "outpoint"
        <*> o .: "sighash"
        <*> (mapM (unmarshalValue ctx) =<< o .:? "redeem")

-- | Sign a transaction by providing the 'SigInput' signing parameters and a
-- list of private keys. The signature is computed deterministically as defined
-- in RFC-6979.
signTx ::
  Network ->
  Ctx ->
  -- | transaction to sign
  Tx ->
  -- | signing parameters, with nesting flag
  [(SigInput, Bool)] ->
  -- | private keys to sign with
  [SecKey] ->
  -- | signed transaction
  Either String Tx
signTx net ctx otx sigis allKeys
  | null ti = Left "signTx: Transaction has no inputs"
  | otherwise = foldM go otx $ findInputIndex ((.outpoint) . fst) sigis ti
  where
    ti = otx.inputs
    go tx (sigi@(SigInput so _ _ _ rdmM, _), i) = do
      keys <- sigKeys ctx so rdmM allKeys
      foldM (\t k -> signInput net ctx t i sigi k) tx keys

-- | Sign a single input in a transaction deterministically (RFC-6979).  The
-- nesting flag only affects the behavior of segwit inputs.
signInput ::
  Network ->
  Ctx ->
  Tx ->
  Int ->
  -- | boolean flag: nest input
  (SigInput, Bool) ->
  PrivateKey ->
  Either String Tx
signInput net ctx tx i (sigIn@(SigInput so val _ _ rdmM), nest) key = do
  let sig = makeSignature net ctx tx i sigIn key
  si <- buildInput net ctx tx i so val rdmM sig $ derivePublicKey ctx key
  w <- updatedWitnessData net ctx tx i so si
  return tx {inputs = nextTxIn so si, witness = w}
  where
    f si TxIn {..} = TxIn {script = marshal (net, ctx) si, ..}
    g so' TxIn {..} = TxIn {script = pkScript so', ..}
    pkScript so' = runPutS . serialize . opPushData $ marshal ctx so'
    nextTxIn so' si
      | isSegwit so' && nest = updateIndex i tx.inputs (g so')
      | isSegwit so' = tx.inputs
      | otherwise = updateIndex i tx.inputs (f si)

-- | Add the witness data of the transaction given segwit parameters for an input.
--
-- @since 0.11.0.0
updatedWitnessData ::
  Network ->
  Ctx ->
  Tx ->
  Int ->
  ScriptOutput ->
  ScriptInput ->
  Either String WitnessData
updatedWitnessData net ctx tx i so si
  | isSegwit so =
      updateWitness . toWitnessStack net ctx =<< calcWitnessProgram net ctx so si
  | otherwise =
      return tx.witness
  where
    updateWitness w
      | null tx.witness = return $ updateIndex i defaultStack (const w)
      | length tx.witness /= n = Left "Invalid number of witness stacks"
      | otherwise = return $ updateIndex i tx.witness (const w)
    defaultStack = replicate n $ toWitnessStack net ctx EmptyWitnessProgram
    n = length tx.inputs

-- | Associate an input index to each value in a list
findInputIndex ::
  -- | extract an outpoint
  (a -> OutPoint) ->
  -- | input list
  [a] ->
  -- | reference list of inputs
  [TxIn] ->
  [(a, Int)]
findInputIndex getOutPoint as ti =
  mapMaybe g $ zip (matchTemplate as ti f) [0 ..]
  where
    f s txin = getOutPoint s == txin.outpoint
    g (Just s, i) = Just (s, i)
    g (Nothing, _) = Nothing

-- | Find from the list of provided private keys which one is required to sign
-- the 'ScriptOutput'.
sigKeys ::
  Ctx ->
  ScriptOutput ->
  Maybe RedeemScript ->
  [SecKey] ->
  Either String [PrivateKey]
sigKeys ctx so rdmM keys =
  case (so, rdmM) of
    (PayPK p, Nothing) ->
      return . map fst . maybeToList $ find ((== p) . snd) zipKeys
    (PayPKHash h, Nothing) -> return $ keyByHash h
    (PayMulSig ps r, Nothing) ->
      return $ map fst $ take r $ filter ((`elem` ps) . snd) zipKeys
    (PayScriptHash _, Just rdm) -> sigKeys ctx rdm Nothing keys
    (PayWitnessPKHash h, _) -> return $ keyByHash h
    (PayWitnessScriptHash _, Just rdm) -> sigKeys ctx rdm Nothing keys
    _ -> Left "sigKeys: Could not decode output script"
  where
    zipKeys =
      [ (prv, pub)
        | k <- keys,
          t <- [True, False],
          let prv = wrapSecKey t k,
          let pub = derivePublicKey ctx prv
      ]
    keyByHash h = fmap fst . maybeToList . findKey h $ zipKeys
    findKey h = find $ (== h) . (.hash160) . pubKeyAddr ctx . snd

-- | Construct an input for a transaction given a signature, public key and data
-- about the previous output.
buildInput ::
  Network ->
  Ctx ->
  -- | transaction where input will be added
  Tx ->
  -- | input index where signature will go
  Int ->
  -- | output script being spent
  ScriptOutput ->
  -- | amount of previous output
  Word64 ->
  -- | redeem script if pay-to-script-hash
  Maybe RedeemScript ->
  TxSignature ->
  PublicKey ->
  Either String ScriptInput
buildInput net ctx tx i so val rdmM sig pub = do
  when (i >= length tx.inputs) $ Left "buildInput: Invalid input index"
  case (so, rdmM) of
    (PayScriptHash _, Just rdm) ->
      buildScriptHashInput rdm
    (PayWitnessScriptHash _, Just rdm) ->
      buildScriptHashInput rdm
    (PayWitnessPKHash _, Nothing) ->
      return . RegularInput $ SpendPKHash sig pub
    (_, Nothing) ->
      buildRegularInput so
    _ -> Left "buildInput: Invalid output/redeem script combination"
  where
    buildRegularInput = \case
      PayPK _ -> return $ RegularInput $ SpendPK sig
      PayPKHash _ -> return $ RegularInput $ SpendPKHash sig pub
      PayMulSig msPubs r -> do
        let mSigs = take r $ catMaybes $ matchTemplate allSigs msPubs f
            allSigs = nub $ sig : parseExistingSigs net ctx tx so i
        return $ RegularInput $ SpendMulSig mSigs
      _ -> Left "buildInput: Invalid output/redeem script combination"
    buildScriptHashInput rdm = do
      inp <- buildRegularInput rdm
      return $ ScriptHashInput inp.get rdm
    f (TxSignature x sh) p =
      verifyHashSig
        ctx
        (makeSigHash net ctx tx i so val sh rdmM)
        x
        p.point
    f TxSignatureEmpty _ = False

-- | Apply heuristics to extract the signatures for a particular input that are
-- embedded in the transaction.
--
-- @since 0.11.0.0
parseExistingSigs :: Network -> Ctx -> Tx -> ScriptOutput -> Int -> [TxSignature]
parseExistingSigs net ctx tx so i = insSigs <> witSigs
  where
    insSigs = case unmarshal (net, ctx) scp of
      Right (ScriptHashInput (SpendMulSig xs) _) -> xs
      Right (RegularInput (SpendMulSig xs)) -> xs
      _ -> []
    scp = (tx.inputs !! i).script
    witSigs
      | not $ isSegwit so = []
      | null tx.witness = []
      | otherwise = rights $ decodeTxSig net ctx <$> (tx.witness !! i)

-- | Produce a structured representation of a deterministic (RFC-6979) signature over an input.
makeSignature :: Network -> Ctx -> Tx -> Int -> SigInput -> PrivateKey -> TxSignature
makeSignature net ctx tx i (SigInput so val _ sh rdmM) key =
  TxSignature (signHash ctx key.key m) sh
  where
    m = makeSigHash net ctx tx i so val sh rdmM

-- | A function which selects the digest algorithm and parameters as appropriate
--
-- @since 0.11.0.0
makeSigHash ::
  Network ->
  Ctx ->
  Tx ->
  Int ->
  ScriptOutput ->
  Word64 ->
  SigHash ->
  Maybe RedeemScript ->
  Hash256
makeSigHash net ctx tx i so val sh rdmM = h net tx (encodeOutput ctx so') val i sh
  where
    so' = case so of
      PayWitnessPKHash h' -> PayPKHash h'
      _ -> fromMaybe so rdmM
    h
      | isSegwit so = txSigHashForkId
      | otherwise = txSigHash