// Copyright (c) 2013-2017 The btcsuite developers
// Copyright (c) 2015-2021 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package txscript

import (
	"bytes"
	"encoding/binary"
	"strings"

	"github.com/decred/dcrd/chaincfg/chainhash"
)

// These are the constants specified for maximums in individual scripts.
const (
	MaxOpsPerScript       = 255  // Max number of non-push operations.
	MaxPubKeysPerMultiSig = 20   // Multisig can't have more sigs than this.
	MaxScriptElementSize  = 2048 // Max bytes pushable to the stack.
)

// IsSmallInt returns whether or not the opcode is considered a small integer,
// which is an OP_0, or OP_1 through OP_16.
//
// NOTE: This function is only valid for version 0 opcodes.  Since the function
// does not accept a script version, the results are undefined for other script
// versions.
func IsSmallInt(op byte) bool {
	return op == OP_0 || (op >= OP_1 && op <= OP_16)
}

// IsPayToScriptHash returns true if the script is in the standard
// pay-to-script-hash (P2SH) format, false otherwise.
//
// WARNING: This function always treats the passed script as version 0.  Great
// care must be taken if introducing a new script version because it is used in
// consensus which, unfortunately as of the time of this writing, does not check
// script versions before determining if the script is a P2SH which means nodes
// on existing rules will analyze new version scripts as if they were version 0.
func IsPayToScriptHash(script []byte) bool {
	return isScriptHashScript(script)
}

// IsPushOnlyScript returns whether or not the passed script only pushes data
// according to the consensus definition of pushing data.
//
// WARNING: This function always treats the passed script as version 0.  Great
// care must be taken if introducing a new script version because it is used in
// consensus which, unfortunately as of the time of this writing, does not check
// script versions before checking if it is a push only script which means nodes
// on existing rules will treat new version scripts as if they were version 0.
func IsPushOnlyScript(script []byte) bool {
	const scriptVersion = 0
	tokenizer := MakeScriptTokenizer(scriptVersion, script)
	for tokenizer.Next() {
		// All opcodes up to OP_16 are data push instructions.
		// NOTE: This does consider OP_RESERVED to be a data push instruction,
		// but execution of OP_RESERVED will fail anyway and matches the
		// behavior required by consensus.
		if tokenizer.Opcode() > OP_16 {
			return false
		}
	}
	return tokenizer.Err() == nil
}

// isStakeOpcode returns whether or not the opcode is one of the stake tagging
// opcodes.
func isStakeOpcode(op byte, isTreasuryEnabled bool) bool {
	if isTreasuryEnabled {
		return (op >= OP_SSTX && op <= OP_SSTXCHANGE) ||
			(op >= OP_TADD && op <= OP_TGEN)
	}
	return op >= OP_SSTX && op <= OP_SSTXCHANGE
}

// ExtractScriptHash extracts the script hash from the passed script if it is a
// standard pay-to-script-hash script.  It will return nil otherwise.
//
// NOTE: This function is only valid for version 0 opcodes.  Since the function
// does not accept a script version, the results are undefined for other script
// versions.
func ExtractScriptHash(script []byte) []byte {
	// A pay-to-script-hash script is of the form:
	//  OP_HASH160 <20-byte scripthash> OP_EQUAL
	if len(script) == 23 &&
		script[0] == OP_HASH160 &&
		script[1] == OP_DATA_20 &&
		script[22] == OP_EQUAL {

		return script[2:22]
	}

	return nil
}

// isScriptHashScript returns whether or not the passed script is a standard
// pay-to-script-hash script.
func isScriptHashScript(script []byte) bool {
	return ExtractScriptHash(script) != nil
}

// isStakeScriptHashScript returns whether or not the passed script is a
// stake-tagged pay-to-script-hash script.
func (vm *Engine) isStakeScriptHashScript(script []byte) bool {
	return len(script) == 24 &&
		isStakeOpcode(script[0], vm.hasFlag(ScriptVerifyTreasury)) &&
		script[1] == OP_HASH160 &&
		script[2] == OP_DATA_20 &&
		script[23] == OP_EQUAL
}

// isAnyKindOfScriptHash returns whether or not the passed script is either a
// regular pay-to-script-hash script or a stake-tagged pay-to-script-hash
// script.
func (vm *Engine) isAnyKindOfScriptHash(script []byte) bool {
	return isScriptHashScript(script) || vm.isStakeScriptHashScript(script)
}

// ContainsStakeOpCodes returns whether or not a public key script contains any
// stake tagging opcodes.
//
// NOTE: This function is only valid for version 0 scripts.  Since the function
// does not accept a script version, the results are undefined for other script
// versions.
func ContainsStakeOpCodes(pkScript []byte, isTreasuryEnabled bool) (bool, error) {
	const scriptVersion = 0
	tokenizer := MakeScriptTokenizer(scriptVersion, pkScript)
	for tokenizer.Next() {
		if isStakeOpcode(tokenizer.Opcode(), isTreasuryEnabled) {
			return true, nil
		}
	}

	return false, tokenizer.Err()
}

// hasP2SHRedeemScriptStakeOpCodes returns an error if the provided public key
// script is a regular pay-to-script-hash or a stake-tagged pay-to-script and,
// when it is, that the redeem script within the provided signature script
// contains stake opcodes.  An error is also returned if the signature script is
// malformed after determining the public key script is one of the
// aforementioned cases.
func (vm *Engine) hasP2SHRedeemScriptStakeOpCodes(version uint16, sigScript, pkScript []byte) error {
	// The only stake scripts currently supported are version 0.
	if version != 0 {
		return nil
	}

	// Nothing further to check if the public key script is not a normal
	// pay-to-script-hash script or one tagged with a stake opcode.
	if !(isScriptHashScript(pkScript) ||
		vm.isStakeScriptHashScript(pkScript)) {
		return nil
	}

	// Extract the redeem script from the signature script.
	redeemScript := finalOpcodeData(version, sigScript)
	if len(redeemScript) == 0 {
		str := "p2sh signature script has no pushed data"
		return scriptError(ErrNotPushOnly, str)
	}

	// Ensure the redeem script does not contain any stake opcodes as their use
	// is prohibited outside of the very specific circumstances permitted by
	// the staking system.
	hasStakeOpCodes, err := ContainsStakeOpCodes(redeemScript,
		vm.hasFlag(ScriptVerifyTreasury))
	if err != nil {
		return err
	}
	if hasStakeOpCodes {
		str := "stake opcodes were found in a p2sh script"
		return scriptError(ErrP2SHStakeOpCodes, str)
	}

	return nil
}

// DisasmString formats a disassembled script for one line printing.  When the
// script fails to parse, the returned string will contain the disassembled
// script up to the point the failure occurred along with the string '[error]'
// appended.  In addition, the reason the script failed to parse is returned
// if the caller wants more information about the failure.
//
// NOTE: This function is only valid for version 0 scripts.  Since the function
// does not accept a script version, the results are undefined for other script
// versions.
func DisasmString(script []byte) (string, error) {
	const scriptVersion = 0

	var disbuf strings.Builder
	tokenizer := MakeScriptTokenizer(scriptVersion, script)
	if tokenizer.Next() {
		disasmOpcode(&disbuf, tokenizer.op, tokenizer.Data(), true)
	}
	for tokenizer.Next() {
		disbuf.WriteByte(' ')
		disasmOpcode(&disbuf, tokenizer.op, tokenizer.Data(), true)
	}
	if tokenizer.Err() != nil {
		if tokenizer.ByteIndex() != 0 {
			disbuf.WriteByte(' ')
		}
		disbuf.WriteString("[error]")
	}
	return disbuf.String(), tokenizer.Err()
}

// isCanonicalPush returns true if the opcode is either not a push instruction
// or the data associated with the push instruction uses the smallest
// instruction to do the job.  False otherwise.
//
// For example, it is possible to push a value of 1 to the stack as "OP_1",
// "OP_DATA_1 0x01", "OP_PUSHDATA1 0x01 0x01", and others, however, the first
// only takes a single byte, while the rest take more.  Only the first is
// considered canonical.
func isCanonicalPush(opcode byte, data []byte) bool {
	dataLen := len(data)
	if opcode > OP_16 {
		return true
	}

	if opcode < OP_PUSHDATA1 && opcode > OP_0 && (dataLen == 1 && data[0] <= 16) {
		return false
	}
	if opcode == OP_PUSHDATA1 && dataLen < OP_PUSHDATA1 {
		return false
	}
	if opcode == OP_PUSHDATA2 && dataLen <= 0xff {
		return false
	}
	if opcode == OP_PUSHDATA4 && dataLen <= 0xffff {
		return false
	}
	return true
}

// removeOpcodeByData will return the script minus any opcodes that perform a
// canonical push of data that contains the passed data to remove.  This
// function assumes it is provided a version 0 script as any future version of
// script should avoid this functionality since it is unnecessary due to the
// signature scripts not being part of the witness-free transaction hash.
//
// WARNING: This will return the passed script unmodified unless a modification
// is necessary in which case the modified script is returned.  This implies
// callers may NOT rely on being able to safely mutate either the passed or
// returned script without potentially modifying the same data.
//
// NOTE: This function is only valid for version 0 scripts.  Since the function
// does not accept a script version, the results are undefined for other script
// versions.
func removeOpcodeByData(script []byte, dataToRemove []byte) []byte {
	// Avoid work when possible.
	if len(script) == 0 || len(dataToRemove) == 0 {
		return script
	}

	// Parse through the script looking for a canonical data push that contains
	// the data to remove.
	const scriptVersion = 0
	var result []byte
	var prevOffset int32
	tokenizer := MakeScriptTokenizer(scriptVersion, script)
	for tokenizer.Next() {
		// In practice, the script will basically never actually contain the
		// data since this function is only used during signature verification
		// to remove the signature itself which would require some incredibly
		// non-standard code to create.
		//
		// Thus, as an optimization, avoid allocating a new script unless there
		// is actually a match that needs to be removed.
		op, data := tokenizer.Opcode(), tokenizer.Data()
		if isCanonicalPush(op, data) && bytes.Contains(data, dataToRemove) {
			if result == nil {
				fullPushLen := tokenizer.ByteIndex() - prevOffset
				result = make([]byte, 0, int32(len(script))-fullPushLen)
				result = append(result, script[0:prevOffset]...)
			}
		} else if result != nil {
			result = append(result, script[prevOffset:tokenizer.ByteIndex()]...)
		}

		prevOffset = tokenizer.ByteIndex()
	}
	if result == nil {
		result = script
	}
	return result
}

// AsSmallInt returns the passed opcode, which MUST be true according to the
// IsSmallInt function, as an integer.
//
// NOTE: This function is only valid for version 0 opcodes.  Since the function
// does not accept a script version, the results are undefined for other script
// versions.
func AsSmallInt(op byte) int {
	if op == OP_0 {
		return 0
	}

	return int(op - (OP_1 - 1))
}

// countSigOpsV0 returns the number of signature operations in the provided
// script up to the point of the first parse failure or the entire script when
// there are no parse failures.  The precise flag attempts to accurately count
// the number of operations for a multisig operation versus using the maximum
// allowed.
//
// WARNING: This function always treats the passed script as version 0.  Great
// care must be taken if introducing a new script version because it is used in
// consensus which, unfortunately as of the time of this writing, does not check
// script versions before counting their signature operations which means nodes
// on existing rules will count new version scripts as if they were version 0.
func countSigOpsV0(script []byte, precise bool, isTreasuryEnabled bool) int {
	const scriptVersion = 0

	numSigOps := 0
	tokenizer := MakeScriptTokenizer(scriptVersion, script)
	prevOp := byte(OP_INVALIDOPCODE)
	for tokenizer.Next() {
		switch tokenizer.Opcode() {
		case OP_TSPEND:
			if isTreasuryEnabled {
				numSigOps++
			}

		case OP_CHECKSIG, OP_CHECKSIGVERIFY, OP_CHECKSIGALT,
			OP_CHECKSIGALTVERIFY:

			numSigOps++

		case OP_CHECKMULTISIG, OP_CHECKMULTISIGVERIFY:
			// Note that OP_0 is treated as the max number of sigops here in
			// precise mode despite it being a valid small integer in order to
			// highly discourage multisigs with zero pubkeys.
			//
			// Also, even though this is referred to as "precise" counting, it's
			// not really precise at all due to the small int opcodes only
			// covering 1 through 16 pubkeys, which means this will count any
			// more than that value (e.g. 17, 18 19) as the maximum number of
			// allowed pubkeys.  This was inherited from bitcoin and is,
			// unfortunately, now part of the consensus rules.  This could be
			// made more correct with a new script version, however, ideally all
			// multisignature operations in new script versions should move to
			// aggregated schemes such as Schnorr instead.
			if precise && prevOp >= OP_1 && prevOp <= OP_16 {
				numSigOps += AsSmallInt(prevOp)
			} else {
				numSigOps += MaxPubKeysPerMultiSig
			}

		default:
			// Not a sigop.
		}

		prevOp = tokenizer.Opcode()
	}

	return numSigOps
}

// GetSigOpCount provides a quick count of the number of signature operations
// in a script. a CHECKSIG operations counts for 1, and a CHECK_MULTISIG for 20.
// If the script fails to parse, then the count up to the point of failure is
// returned.
//
// WARNING: This function always treats the passed script as version 0.  Great
// care must be taken if introducing a new script version because it is used in
// consensus which, unfortunately as of the time of this writing, does not check
// script versions before counting their signature operations which means nodes
// on existing rules will count new version scripts as if they were version 0.
func GetSigOpCount(script []byte, isTreasuryEnabled bool) int {
	return countSigOpsV0(script, false, isTreasuryEnabled)
}

// finalOpcodeData returns the data associated with the final opcode in the
// script.  It will return nil if the script fails to parse.
func finalOpcodeData(scriptVersion uint16, script []byte) []byte {
	// Avoid unnecessary work.
	if len(script) == 0 {
		return nil
	}

	var data []byte
	tokenizer := MakeScriptTokenizer(scriptVersion, script)
	for tokenizer.Next() {
		data = tokenizer.Data()
	}
	if tokenizer.Err() != nil {
		return nil
	}
	return data
}

// GetPreciseSigOpCount returns the number of signature operations in
// scriptPubKey.  If bip16 is true then scriptSig may be searched for the
// Pay-To-Script-Hash script in order to find the precise number of signature
// operations in the transaction.  If the script fails to parse, then the count
// up to the point of failure is returned.
//
// WARNING: This function always treats the passed script as version 0.  Great
// care must be taken if introducing a new script version because it is used in
// consensus which, unfortunately as of the time of this writing, does not check
// script versions before counting their signature operations which means nodes
// on existing rules will count new version scripts as if they were version 0.
func GetPreciseSigOpCount(scriptSig, scriptPubKey []byte, isTreasuryEnabled bool) int {
	const scriptVersion = 0

	// Treat non P2SH transactions as normal.  Note that signature operation
	// counting includes all operations up to the first parse failure.
	if !isScriptHashScript(scriptPubKey) {
		return countSigOpsV0(scriptPubKey, true, isTreasuryEnabled)
	}

	// The signature script must only push data to the stack for P2SH to be
	// a valid pair, so the signature operation count is 0 when that is not
	// the case.
	if len(scriptSig) == 0 || !IsPushOnlyScript(scriptSig) {
		return 0
	}

	// The P2SH script is the last item the signature script pushes to the
	// stack.  When the script is empty, there are no signature operations.
	//
	// Notice that signature scripts that fail to fully parse count as 0
	// signature operations unlike public key and redeem scripts.
	redeemScript := finalOpcodeData(scriptVersion, scriptSig)
	if len(redeemScript) == 0 {
		return 0
	}

	// Return the more precise sigops count for the redeem script.  Note that
	// signature operation counting includes all operations up to the first
	// parse failure.
	return countSigOpsV0(redeemScript, true, isTreasuryEnabled)
}

// checkScriptParses returns an error if the provided script fails to parse.
func checkScriptParses(scriptVersion uint16, script []byte) error {
	tokenizer := MakeScriptTokenizer(scriptVersion, script)
	for tokenizer.Next() {
		// Nothing to do.
	}
	return tokenizer.Err()
}

// IsUnspendable returns whether the passed public key script is unspendable, or
// guaranteed to fail at execution.  This allows inputs to be pruned instantly
// when entering the UTXO set. In Decred, all zero value outputs are unspendable.
//
// NOTE: This function is only valid for version 0 scripts.  Since the function
// does not accept a script version, the results are undefined for other script
// versions.
func IsUnspendable(amount int64, pkScript []byte) bool {
	// The script is unspendable if starts with OP_RETURN or is guaranteed to
	// fail at execution due to being larger than the max allowed script size.
	if amount == 0 || len(pkScript) > MaxScriptSize || len(pkScript) > 0 &&
		pkScript[0] == OP_RETURN {

		return true
	}

	// The script is unspendable if it is guaranteed to fail at execution.
	const scriptVersion = 0
	return checkScriptParses(scriptVersion, pkScript) != nil
}

// GenerateSSGenBlockRef generates a block reference script for the given block
// hash and height which a block votes on.  The script is for use in stake vote
// transactions.
func GenerateSSGenBlockRef(blockHash chainhash.Hash, height uint32) ([]byte, error) {
	// Serialize the block hash and height
	brBytes := make([]byte, 32+4)
	copy(brBytes[0:32], blockHash[:])
	binary.LittleEndian.PutUint32(brBytes[32:36], height)

	return NewScriptBuilder().AddOp(OP_RETURN).AddData(brBytes).Script()
}

// GenerateSSGenVotes generates a vote script for the given vote bits.  The
// script is for use in stake vote transactions.
func GenerateSSGenVotes(votebits uint16) ([]byte, error) {
	// Serialize the votebits
	vbBytes := make([]byte, 2)
	binary.LittleEndian.PutUint16(vbBytes, votebits)

	return NewScriptBuilder().AddOp(OP_RETURN).AddData(vbBytes).Script()
}
