content: {-# LANGUAGE DeriveAnyClass #-}
  {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE DuplicateRecordFields #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE ImportQualifiedPost #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE OverloadedRecordDot #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE NoFieldSelectors #-}
  
  -- |
  -- Module      : Haskoin.Script.Standard
  -- Copyright   : No rights reserved
  -- License     : MIT
  -- Maintainer  : jprupp@protonmail.ch
  -- Stability   : experimental
  -- Portability : POSIX
  --
  -- Standard scripts like pay-to-public-key, pay-to-public-key-hash,
  -- pay-to-script-hash, pay-to-multisig and corresponding SegWit variants.
  module Haskoin.Script.Standard
    ( -- * Standard Script Outputs
      ScriptOutput (..),
      RedeemScript,
      isPayPK,
      isPayPKHash,
      isPayMulSig,
      isPayScriptHash,
      isPayWitness,
      isPayWitnessPKHash,
      isPayWitnessScriptHash,
      isDataCarrier,
      encodeOutput,
      decodeOutput,
      toP2SH,
      toP2WSH,
      sortMulSig,
  
      -- * Standard Script Inputs
      ScriptInput (..),
      SimpleInput (..),
      encodeInput,
      decodeInput,
      isSpendPK,
      isSpendPKHash,
      isSpendMulSig,
      isScriptHashInput,
    )
  where
  
  import Control.Applicative ((<|>))
  import Control.DeepSeq
  import Control.Monad (guard, liftM2, (<=<))
  import Crypto.Secp256k1
  import Data.Aeson (ToJSON (..), Value (..), withText)
  import Data.Aeson.Encoding (Encoding, text)
  import Data.Aeson.Types (Parser)
  import Data.ByteString (ByteString)
  import Data.ByteString qualified as B
  import Data.Bytes.Get
  import Data.Bytes.Put
  import Data.Bytes.Serial
  import Data.Function (on)
  import Data.Hashable
  import Data.List (sortBy)
  import Data.Maybe (fromJust, isJust)
  import Data.Word (Word8)
  import GHC.Generics (Generic)
  import Haskoin.Crypto.Hash
  import Haskoin.Crypto.Keys.Common
  import Haskoin.Network.Data
  import Haskoin.Script.Common
  import Haskoin.Script.SigHash
  import Haskoin.Util
  
  -- | Data type describing standard transaction output scripts. Output scripts
  -- provide the conditions that must be fulfilled for someone to spend the funds
  -- in a transaction output.
  data ScriptOutput
    = -- | pay to public key
      PayPK {key :: !PublicKey}
    | -- | pay to public key hash
      PayPKHash {hash160 :: !Hash160}
    | -- | multisig
      PayMulSig
        { keys :: ![PublicKey],
          required :: !Int
        }
    | -- | pay to a script hash
      PayScriptHash {hash160 :: !Hash160}
    | -- | pay to witness public key hash
      PayWitnessPKHash {hash160 :: !Hash160}
    | -- | pay to witness script hash
      PayWitnessScriptHash {hash256 :: !Hash256}
    | -- | another pay to witness address
      PayWitness
        { version :: !Word8,
          bytes :: !ByteString
        }
    | -- | provably unspendable data carrier
      DataCarrier {bytes :: !ByteString}
    deriving (Eq, Show, Read, Generic, NFData)
  
  instance MarshalJSON Ctx ScriptOutput where
    unmarshalValue ctx =
      withText "ScriptOutput" $ \t ->
        case decodeHex t of
          Nothing -> fail "Could not decode hex script"
          Just bs -> either fail return $ unmarshal ctx bs
  
    marshalValue ctx = String . encodeHex . marshal ctx
  
    marshalEncoding ctx = hexEncoding . runPutL . marshalPut ctx
  
  -- | Is script a pay-to-public-key output?
  isPayPK :: ScriptOutput -> Bool
  isPayPK (PayPK _) = True
  isPayPK _ = False
  
  -- | Is script a pay-to-pub-key-hash output?
  isPayPKHash :: ScriptOutput -> Bool
  isPayPKHash (PayPKHash _) = True
  isPayPKHash _ = False
  
  -- | Is script a pay-to-multi-sig output?
  isPayMulSig :: ScriptOutput -> Bool
  isPayMulSig (PayMulSig _ _) = True
  isPayMulSig _ = False
  
  -- | Is script a pay-to-script-hash output?
  isPayScriptHash :: ScriptOutput -> Bool
  isPayScriptHash (PayScriptHash _) = True
  isPayScriptHash _ = False
  
  -- | Is script a pay-to-witness-pub-key-hash output?
  isPayWitnessPKHash :: ScriptOutput -> Bool
  isPayWitnessPKHash (PayWitnessPKHash _) = True
  isPayWitnessPKHash _ = False
  
  -- | Is script a pay-to-witness-script-hash output?
  isPayWitnessScriptHash :: ScriptOutput -> Bool
  isPayWitnessScriptHash (PayWitnessScriptHash _) = True
  isPayWitnessScriptHash _ = False
  
  -- | Is script paying to a different type of witness address?
  isPayWitness :: ScriptOutput -> Bool
  isPayWitness (PayWitness _ _) = True
  isPayWitness _ = False
  
  -- | Is script a data carrier output?
  isDataCarrier :: ScriptOutput -> Bool
  isDataCarrier (DataCarrier _) = True
  isDataCarrier _ = False
  
  -- | Tries to decode a 'ScriptOutput' from a 'Script'. This can fail if the
  -- script is not recognized as any of the standard output types.
  decodeOutput :: Ctx -> Script -> Either String ScriptOutput
  decodeOutput ctx s = case s.ops of
    -- Pay to PubKey
    [OP_PUSHDATA bs _, OP_CHECKSIG] ->
      PayPK <$> unmarshal ctx bs
    -- Pay to PubKey Hash
    [OP_DUP, OP_HASH160, OP_PUSHDATA bs _, OP_EQUALVERIFY, OP_CHECKSIG] ->
      PayPKHash <$> runGetS deserialize bs
    -- Pay to Script Hash
    [OP_HASH160, OP_PUSHDATA bs _, OP_EQUAL] ->
      PayScriptHash <$> runGetS deserialize bs
    -- Pay to Witness
    [OP_0, OP_PUSHDATA bs OPCODE]
      | B.length bs == 20 ->
          PayWitnessPKHash <$> runGetS deserialize bs
      | B.length bs == 32 ->
          PayWitnessScriptHash <$> runGetS deserialize bs
      | B.length bs /= 20 && B.length bs /= 32 ->
          Left
            "decodeOutput: invalid version 0 segwit \
            \(must be 20 or 32 bytes)"
    -- Other Witness
    [ver, OP_PUSHDATA bs _]
      | Just wv <- opWitnessVersion ver,
        B.length bs >= 2,
        B.length bs <= 40 ->
          Right $ PayWitness wv bs
    -- Provably unspendable data carrier output
    [OP_RETURN, OP_PUSHDATA bs _] -> Right $ DataCarrier bs
    -- Pay to MultiSig Keys
    _ -> case matchPayMulSig ctx s of
      Right x -> return x
      Left _ -> Left "decodeOutput: Non-standard output"
  
  witnessVersionOp :: Word8 -> Maybe ScriptOp
  witnessVersionOp 0 = Just OP_0
  witnessVersionOp 1 = Just OP_1
  witnessVersionOp 2 = Just OP_2
  witnessVersionOp 3 = Just OP_3
  witnessVersionOp 4 = Just OP_4
  witnessVersionOp 5 = Just OP_5
  witnessVersionOp 6 = Just OP_6
  witnessVersionOp 7 = Just OP_7
  witnessVersionOp 8 = Just OP_8
  witnessVersionOp 9 = Just OP_9
  witnessVersionOp 10 = Just OP_10
  witnessVersionOp 11 = Just OP_11
  witnessVersionOp 12 = Just OP_12
  witnessVersionOp 13 = Just OP_13
  witnessVersionOp 14 = Just OP_14
  witnessVersionOp 15 = Just OP_15
  witnessVersionOp 16 = Just OP_16
  witnessVersionOp _ = Nothing
  
  opWitnessVersion :: ScriptOp -> Maybe Word8
  opWitnessVersion OP_0 = Just 0
  opWitnessVersion OP_1 = Just 1
  opWitnessVersion OP_2 = Just 2
  opWitnessVersion OP_3 = Just 3
  opWitnessVersion OP_4 = Just 4
  opWitnessVersion OP_5 = Just 5
  opWitnessVersion OP_6 = Just 6
  opWitnessVersion OP_7 = Just 7
  opWitnessVersion OP_8 = Just 8
  opWitnessVersion OP_9 = Just 9
  opWitnessVersion OP_10 = Just 10
  opWitnessVersion OP_11 = Just 11
  opWitnessVersion OP_12 = Just 12
  opWitnessVersion OP_13 = Just 13
  opWitnessVersion OP_14 = Just 14
  opWitnessVersion OP_15 = Just 15
  opWitnessVersion OP_16 = Just 16
  opWitnessVersion _ = Nothing
  
  -- | Computes a 'Script' from a standard 'ScriptOutput'.
  encodeOutput :: Ctx -> ScriptOutput -> Script
  encodeOutput ctx s = Script $ case s of
    -- Pay to PubKey
    (PayPK k) -> [opPushData $ marshal ctx k, OP_CHECKSIG]
    -- Pay to PubKey Hash Address
    (PayPKHash h) ->
      [ OP_DUP,
        OP_HASH160,
        opPushData $ runPutS $ serialize h,
        OP_EQUALVERIFY,
        OP_CHECKSIG
      ]
    -- Pay to MultiSig Keys
    (PayMulSig ps r)
      | r <= length ps ->
          let opM = intToScriptOp r
              opN = intToScriptOp $ length ps
              keys = map (opPushData . marshal ctx) ps
           in opM : keys ++ [opN, OP_CHECKMULTISIG]
      | otherwise -> error "encodeOutput: PayMulSig r must be <= than pkeys"
    -- Pay to Script Hash Address
    (PayScriptHash h) ->
      [OP_HASH160, opPushData $ runPutS $ serialize h, OP_EQUAL]
    -- Pay to Witness PubKey Hash Address
    (PayWitnessPKHash h) ->
      [OP_0, opPushData $ runPutS $ serialize h]
    (PayWitnessScriptHash h) ->
      [OP_0, opPushData $ runPutS $ serialize h]
    (PayWitness v h) ->
      [ case witnessVersionOp v of
          Nothing -> error "encodeOutput: invalid witness version"
          Just c -> c,
        opPushData h
      ]
    -- Provably unspendable output
    (DataCarrier d) -> [OP_RETURN, opPushData d]
  
  instance Marshal Ctx ScriptOutput where
    marshalGet ctx = do
      script <- deserialize
      case decodeOutput ctx script of
        Left e -> fail e
        Right o -> return o
    marshalPut ctx = serialize . encodeOutput ctx
  
  -- | Encode script as pay-to-script-hash script
  toP2SH :: Script -> ScriptOutput
  toP2SH = PayScriptHash . addressHash . runPutS . serialize
  
  -- | Encode script as a pay-to-witness-script-hash script
  toP2WSH :: Script -> ScriptOutput
  toP2WSH = PayWitnessScriptHash . sha256 . runPutS . serialize
  
  -- | Match @[OP_N, PubKey1, ..., PubKeyM, OP_M, OP_CHECKMULTISIG]@
  matchPayMulSig :: Ctx -> Script -> Either String ScriptOutput
  matchPayMulSig ctx (Script ops) = case splitAt (length ops - 2) ops of
    (m : xs, [n, OP_CHECKMULTISIG]) -> do
      (intM, intN) <- liftM2 (,) (scriptOpToInt m) (scriptOpToInt n)
      if intM <= intN && length xs == intN
        then liftM2 PayMulSig (go xs) (return intM)
        else Left "matchPayMulSig: Invalid M or N parameters"
    _ -> Left "matchPayMulSig: script did not match output template"
    where
      go (OP_PUSHDATA bs _ : xs) =
        liftM2 (:) (unmarshal ctx bs) (go xs)
      go [] =
        Right []
      go _ =
        Left "matchPayMulSig: invalid multisig opcode"
  
  -- | Sort the public keys of a multisig output in ascending order by comparing
  -- their compressed serialized representations. Refer to BIP-67.
  sortMulSig :: Ctx -> ScriptOutput -> ScriptOutput
  sortMulSig ctx out = case out of
    PayMulSig keys r ->
      PayMulSig
        (sortBy (compare `on` marshal ctx) keys)
        r
    _ -> error "Can only call orderMulSig on PayMulSig scripts"
  
  -- | Data type describing standard transaction input scripts. Input scripts
  -- provide the signing data required to unlock the coins of the output they are
  -- trying to spend, except in pay-to-witness-public-key-hash and
  -- pay-to-script-hash transactions.
  data SimpleInput
    = SpendPK
        { -- | transaction signature
          signature :: !TxSignature
        }
    | SpendPKHash
        { -- | embedded signature
          signature :: !TxSignature,
          -- | public key
          key :: !PublicKey
        }
    | SpendMulSig
        { -- | list of signatures
          signatures :: ![TxSignature]
        }
    deriving (Eq, Show, Read, Generic, NFData)
  
  -- | Returns true if the input script is spending from a pay-to-public-key
  -- output.
  isSpendPK :: ScriptInput -> Bool
  isSpendPK (RegularInput (SpendPK _)) = True
  isSpendPK _ = False
  
  -- | Returns true if the input script is spending from a pay-to-public-key-hash
  -- output.
  isSpendPKHash :: ScriptInput -> Bool
  isSpendPKHash (RegularInput (SpendPKHash _ _)) = True
  isSpendPKHash _ = False
  
  -- | Returns true if the input script is spending a multisig output.
  isSpendMulSig :: ScriptInput -> Bool
  isSpendMulSig (RegularInput (SpendMulSig _)) = True
  isSpendMulSig _ = False
  
  -- | Returns true if the input script is spending a pay-to-script-hash output.
  isScriptHashInput :: ScriptInput -> Bool
  isScriptHashInput (ScriptHashInput _ _) = True
  isScriptHashInput _ = False
  
  -- | A redeem script is the output script serialized into the spending input
  -- script. It must be included in inputs that spend pay-to-script-hash outputs.
  type RedeemScript = ScriptOutput
  
  -- | Standard input script high-level representation.
  data ScriptInput
    = RegularInput
        { -- | get wrapped simple input
          get :: !SimpleInput
        }
    | ScriptHashInput
        { -- | get simple input associated with redeem script
          get :: !SimpleInput,
          -- | redeem script
          redeem :: !RedeemScript
        }
    deriving (Show, Read, Eq, Generic, NFData)
  
  -- | Heuristic to decode an input script into one of the standard types.
  decodeSimpleInput :: Network -> Ctx -> Script -> Either String SimpleInput
  decodeSimpleInput net ctx (Script ops) =
    maybeToEither errMsg $ matchPK ops <|> matchPKHash ops <|> matchMulSig ops
    where
      matchPK [op] = SpendPK <$> f op
      matchPK _ = Nothing
      matchPKHash [op, OP_PUSHDATA pub _] =
        SpendPKHash <$> f op <*> eitherToMaybe (unmarshal ctx pub)
      matchPKHash _ = Nothing
      matchMulSig (x : xs) = do
        guard $ x == OP_0
        SpendMulSig <$> mapM f xs
      matchMulSig _ = Nothing
      f OP_0 = return TxSignatureEmpty
      f (OP_PUSHDATA "" OPCODE) = f OP_0
      f (OP_PUSHDATA bs _) = eitherToMaybe $ decodeTxSig net ctx bs
      f _ = Nothing
      errMsg = "decodeInput: Could not decode script input"
  
  -- | Heuristic to decode a 'ScriptInput' from a 'Script'. This function fails if
  -- the script can not be parsed as a standard script input.
  decodeInput :: Network -> Ctx -> Script -> Either String ScriptInput
  decodeInput net ctx s@(Script ops) =
    maybeToEither errMsg $ matchSimpleInput <|> matchPayScriptHash
    where
      matchSimpleInput =
        RegularInput <$> eitherToMaybe (decodeSimpleInput net ctx s)
      matchPayScriptHash =
        case splitAt (length s.ops - 1) ops of
          (is, [OP_PUSHDATA bs _]) -> do
            rdm <- eitherToMaybe $ unmarshal ctx bs
            inp <- eitherToMaybe $ decodeSimpleInput net ctx $ Script is
            return $ ScriptHashInput inp rdm
          _ -> Nothing
      errMsg = "decodeInput: Could not decode script input"
  
  instance Marshal (Network, Ctx) ScriptInput where
    marshalGet (net, ctx) =
      deserialize >>= either fail return . decodeInput net ctx
  
    marshalPut (net, ctx) =
      serialize . encodeInput net ctx
  
  -- | Encode a standard input into a script.
  encodeInput :: Network -> Ctx -> ScriptInput -> Script
  encodeInput net ctx s = case s of
    RegularInput ri -> encodeSimpleInput net ctx ri
    ScriptHashInput i o ->
      Script $ (encodeSimpleInput net ctx i).ops ++ [opPushData $ marshal ctx o]
  
  -- | Encode a standard 'SimpleInput' into opcodes as an input 'Script'.
  encodeSimpleInput :: Network -> Ctx -> SimpleInput -> Script
  encodeSimpleInput net ctx s =
    Script $
      case s of
        SpendPK ts -> [f ts]
        SpendPKHash ts p -> [f ts, opPushData $ marshal ctx p]
        SpendMulSig xs -> OP_0 : map f xs
    where
      f TxSignatureEmpty = OP_0
      f ts = opPushData $ encodeTxSig net ctx ts
  