package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"strconv"

	"github.com/btcsuite/btcd/btcec/v2/schnorr"
	"github.com/btcsuite/btcd/btcutil"
	"github.com/btcsuite/btcd/btcutil/hdkeychain"
	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/mempool"
	"github.com/btcsuite/btcd/txscript"
	"github.com/btcsuite/btcd/wire"
	"github.com/decred/dcrd/dcrec/secp256k1/v4"
	"github.com/lightninglabs/chantools/lnd"
	"github.com/lightningnetwork/lnd/input"
	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
	"github.com/spf13/cobra"
)

type doubleSpendInputs struct {
	APIURL         string
	InputOutpoints []string
	Publish        bool
	SweepAddr      string
	FeeRate        uint32
	RecoveryWindow uint32

	rootKey *rootKey
	cmd     *cobra.Command
}

func newDoubleSpendInputsCommand() *cobra.Command {
	cc := &doubleSpendInputs{}
	cc.cmd = &cobra.Command{
		Use:   "doublespendinputs",
		Short: "Replace a transaction by double spending its input",
		Long: `Tries to double spend the given inputs by deriving the
private for the address and sweeping the funds to the given address. This can
only be used with inputs that belong to an lnd wallet.`,
		Example: `chantools doublespendinputs \
	--inputoutpoints xxxxxxxxx:y,xxxxxxxxx:y \
	--sweepaddr bc1q..... \
	--feerate 10 \
	--publish`,
		RunE: cc.Execute,
	}
	cc.cmd.Flags().StringVar(
		&cc.APIURL, "apiurl", defaultAPIURL, "API URL to use (must "+
			"be esplora compatible)",
	)
	cc.cmd.Flags().StringSliceVar(
		&cc.InputOutpoints, "inputoutpoints", []string{},
		"list of outpoints to double spend in the format txid:vout",
	)
	cc.cmd.Flags().StringVar(
		&cc.SweepAddr, "sweepaddr", "", "address to recover the funds "+
			"to; specify '"+lnd.AddressDeriveFromWallet+"' to "+
			"derive a new address from the seed automatically",
	)
	cc.cmd.Flags().Uint32Var(
		&cc.FeeRate, "feerate", defaultFeeSatPerVByte, "fee rate to "+
			"use for the sweep transaction in sat/vByte",
	)
	cc.cmd.Flags().Uint32Var(
		&cc.RecoveryWindow, "recoverywindow", defaultRecoveryWindow,
		"number of keys to scan per internal/external branch; output "+
			"will consist of double this amount of keys",
	)
	cc.cmd.Flags().BoolVar(
		&cc.Publish, "publish", false, "publish replacement TX to "+
			"the chain API instead of just printing the TX",
	)

	cc.rootKey = newRootKey(cc.cmd, "deriving the input keys")

	return cc.cmd
}

func (c *doubleSpendInputs) Execute(_ *cobra.Command, _ []string) error {
	extendedKey, err := c.rootKey.read()
	if err != nil {
		return fmt.Errorf("error reading root key: %w", err)
	}

	// Make sure sweep addr is set.
	err = lnd.CheckAddress(
		c.SweepAddr, chainParams, true, "sweep", lnd.AddrTypeP2WKH,
		lnd.AddrTypeP2TR,
	)
	if err != nil {
		return err
	}

	// Make sure we have at least one input.
	if len(c.InputOutpoints) == 0 {
		return fmt.Errorf("inputoutpoints are required")
	}

	api := newExplorerAPI(c.APIURL)

	addresses := make([]btcutil.Address, 0, len(c.InputOutpoints))
	outpoints := make([]*wire.OutPoint, 0, len(c.InputOutpoints))
	privKeys := make([]*secp256k1.PrivateKey, 0, len(c.InputOutpoints))

	// Get the addresses for the inputs.
	for _, inputOutpoint := range c.InputOutpoints {
		addrString, err := api.Address(inputOutpoint)
		if err != nil {
			return err
		}

		addr, err := btcutil.DecodeAddress(addrString, chainParams)
		if err != nil {
			return err
		}

		addresses = append(addresses, addr)

		txHash, err := chainhash.NewHashFromStr(inputOutpoint[:64])
		if err != nil {
			return err
		}

		vout, err := strconv.Atoi(inputOutpoint[65:])
		if err != nil {
			return err
		}
		outpoint := wire.NewOutPoint(txHash, uint32(vout))

		outpoints = append(outpoints, outpoint)
	}

	// Create the paths for the addresses.
	p2wkhPath, err := lnd.ParsePath(lnd.WalletDefaultDerivationPath)
	if err != nil {
		return err
	}

	p2trPath, err := lnd.ParsePath(lnd.WalletBIP86DerivationPath)
	if err != nil {
		return err
	}

	// Start with the txweight estimator.
	var estimator input.TxWeightEstimator
	sweepScript, err := lnd.PrepareWalletAddress(
		c.SweepAddr, chainParams, &estimator, extendedKey, "sweep",
	)
	if err != nil {
		return err
	}

	// Find the key for the given addresses and add their
	// output weight to the tx estimator.
	for _, addr := range addresses {
		var key *hdkeychain.ExtendedKey
		switch addr.(type) {
		case *btcutil.AddressWitnessPubKeyHash:
			key, err = iterateOverPath(
				extendedKey, addr, p2wkhPath, c.RecoveryWindow,
			)
			if err != nil {
				return err
			}

			estimator.AddP2WKHInput()

		case *btcutil.AddressTaproot:
			key, err = iterateOverPath(
				extendedKey, addr, p2trPath, c.RecoveryWindow,
			)
			if err != nil {
				return err
			}

			estimator.AddTaprootKeySpendInput(
				txscript.SigHashDefault,
			)

		default:
			return fmt.Errorf("address type %T not supported", addr)
		}

		// Get the private key.
		privKey, err := key.ECPrivKey()
		if err != nil {
			return err
		}

		privKeys = append(privKeys, privKey)
	}

	// Now that we have the keys, we can create the transaction.
	prevOuts := make(map[wire.OutPoint]*wire.TxOut)

	// Next get the full value of the inputs.
	var totalInput btcutil.Amount
	for _, outpoint := range outpoints {
		// Get the transaction.
		tx, err := api.Transaction(outpoint.Hash.String())
		if err != nil {
			return err
		}

		value := tx.Vout[outpoint.Index].Value

		// Get the output index.
		totalInput += btcutil.Amount(value)

		scriptPubkey, err := hex.DecodeString(
			tx.Vout[outpoint.Index].ScriptPubkey,
		)
		if err != nil {
			return err
		}

		// Add the output to the map.
		prevOuts[*outpoint] = &wire.TxOut{
			Value:    int64(value),
			PkScript: scriptPubkey,
		}
	}

	// Calculate the fee.
	feeRateKWeight := chainfee.SatPerKVByte(1000 * c.FeeRate).FeePerKWeight()
	totalFee := feeRateKWeight.FeeForWeight(int64(estimator.Weight()))

	// Create the transaction.
	tx := wire.NewMsgTx(2)

	// Add the inputs.
	for _, outpoint := range outpoints {
		tx.AddTxIn(&wire.TxIn{
			PreviousOutPoint: *outpoint,
			Sequence:         mempool.MaxRBFSequence,
		})
	}

	tx.AddTxOut(wire.NewTxOut(int64(totalInput-totalFee), sweepScript))

	// Calculate the signature hash.
	prevOutFetcher := txscript.NewMultiPrevOutFetcher(prevOuts)
	sigHashes := txscript.NewTxSigHashes(tx, prevOutFetcher)

	// Sign the inputs depending on the address type.
	for i, outpoint := range outpoints {
		switch addresses[i].(type) {
		case *btcutil.AddressWitnessPubKeyHash:
			witness, err := txscript.WitnessSignature(
				tx, sigHashes, i, prevOuts[*outpoint].Value,
				prevOuts[*outpoint].PkScript,
				txscript.SigHashAll, privKeys[i], true,
			)
			if err != nil {
				return err
			}

			tx.TxIn[i].Witness = witness

		case *btcutil.AddressTaproot:
			rawTxSig, err := txscript.RawTxInTaprootSignature(
				tx, sigHashes, i,
				prevOuts[*outpoint].Value,
				prevOuts[*outpoint].PkScript,
				[]byte{}, txscript.SigHashDefault, privKeys[i],
			)
			if err != nil {
				return err
			}

			tx.TxIn[i].Witness = wire.TxWitness{
				rawTxSig,
			}

		default:
			return fmt.Errorf("address type %T not supported",
				addresses[i])
		}
	}

	// Serialize the transaction.
	var txBuf bytes.Buffer
	if err := tx.Serialize(&txBuf); err != nil {
		return err
	}

	// Print the transaction.
	fmt.Printf("Sweeping transaction:\n%x\n", txBuf.Bytes())

	// Publish the transaction.
	if c.Publish {
		txid, err := api.PublishTx(hex.EncodeToString(txBuf.Bytes()))
		if err != nil {
			return err
		}

		fmt.Printf("Published transaction with txid %s\n", txid)
	}

	return nil
}

// iterateOverPath iterates over the given key path and tries to find the
// private key that corresponds to the given address.
func iterateOverPath(baseKey *hdkeychain.ExtendedKey, addr btcutil.Address,
	path []uint32, maxTries uint32) (*hdkeychain.ExtendedKey, error) {

	for i := uint32(0); i < maxTries; i++ {
		// Check for both the external and internal branch.
		for _, branch := range []uint32{0, 1} {
			// Create the path to derive the key.
			addrPath := append(path, branch, i) //nolint:gocritic

			// Derive the key.
			derivedKey, err := lnd.DeriveChildren(baseKey, addrPath)
			if err != nil {
				return nil, err
			}

			var address btcutil.Address
			switch addr.(type) {
			case *btcutil.AddressWitnessPubKeyHash:
				// Get the address for the derived key.
				derivedAddr, err := derivedKey.Address(chainParams)
				if err != nil {
					return nil, err
				}

				address, err = btcutil.NewAddressWitnessPubKeyHash(
					derivedAddr.ScriptAddress(), chainParams,
				)
				if err != nil {
					return nil, err
				}

			case *btcutil.AddressTaproot:

				pubkey, err := derivedKey.ECPubKey()
				if err != nil {
					return nil, err
				}

				pubkey = txscript.ComputeTaprootKeyNoScript(pubkey)

				address, err = btcutil.NewAddressTaproot(
					schnorr.SerializePubKey(pubkey), chainParams,
				)
				if err != nil {
					return nil, err
				}
			}

			// Compare the addresses.
			if address.String() == addr.String() {
				return derivedKey, nil
			}
		}
	}

	return nil, fmt.Errorf("could not find key for address %s", addr.String())
}