package tree

import (
	"bytes"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"runtime"
	"sync"

	"github.com/arkade-os/arkd/pkg/ark-lib/script"
	"github.com/arkade-os/arkd/pkg/ark-lib/txutils"
	"github.com/btcsuite/btcd/btcec/v2"
	"github.com/btcsuite/btcd/btcec/v2/schnorr"
	"github.com/btcsuite/btcd/btcec/v2/schnorr/musig2"
	"github.com/btcsuite/btcd/btcutil/psbt"
	"github.com/btcsuite/btcd/txscript"
	"github.com/btcsuite/btcd/wire"
)

var (
	ErrMissingVtxoTree = errors.New("missing vtxo tree")
)

type Musig2Nonce struct {
	PubNonce [66]byte
}

// TreeNonces is a map of txid to public nonces only
// it implements json.Marshaler and json.Unmarshaler
type TreeNonces map[string]*Musig2Nonce // txid -> public nonces only

func NewTreeNonces(noncesMap map[string]string) (TreeNonces, error) {
	nonces := make(TreeNonces)
	for txid, nonce := range noncesMap {
		nonceBytes, err := hex.DecodeString(nonce)
		if err != nil {
			return nil, err
		}

		if len(nonceBytes) != 66 {
			return nil, fmt.Errorf("expected nonce to be 66 bytes, got %d", len(nonceBytes))
		}

		nonces[txid] = &Musig2Nonce{
			PubNonce: [66]byte(nonceBytes),
		}
	}
	return nonces, nil
}

func (n TreeNonces) ToMap() map[string]string {
	m := make(map[string]string)
	for txid, nonce := range n {
		m[txid] = hex.EncodeToString(nonce.PubNonce[:])
	}
	return m
}

func (n TreeNonces) MarshalJSON() ([]byte, error) {
	mapObject := make(map[string]string)
	for txid, nonce := range n {
		mapObject[txid] = hex.EncodeToString(nonce.PubNonce[:])
	}

	return json.Marshal(mapObject)
}

func (n *TreeNonces) UnmarshalJSON(data []byte) error {
	mapObject := make(map[string]string)
	if err := json.Unmarshal(data, &mapObject); err != nil {
		return err
	}

	*n = make(TreeNonces)

	for txid, nonce := range mapObject {
		nonceBytes, err := hex.DecodeString(nonce)
		if err != nil {
			return err
		}

		if len(nonceBytes) != 66 {
			return fmt.Errorf("expected nonce to be 66 bytes, got %d", len(nonceBytes))
		}

		(*n)[txid] = &Musig2Nonce{
			PubNonce: [66]byte(nonceBytes),
		}
	}

	return nil
}

// TreePartialSigs is a map of txid to partial signature
// it implements json.Marshaler and json.Unmarshaler
type TreePartialSigs map[string]*musig2.PartialSignature // txid -> partial signature

func NewTreePartialSigs(sigsMap map[string]string) (TreePartialSigs, error) {
	sigs := make(TreePartialSigs)
	for txid, sig := range sigsMap {
		sigBytes, err := hex.DecodeString(sig)
		if err != nil {
			return nil, err
		}

		sig := &musig2.PartialSignature{}
		if err := sig.Decode(bytes.NewReader(sigBytes)); err != nil {
			return nil, err
		}

		sigs[txid] = sig
	}
	return sigs, nil
}

func (s TreePartialSigs) ToMap() (map[string]string, error) {
	m := make(map[string]string)
	for txid, sig := range s {
		var sigBytes bytes.Buffer
		if err := sig.Encode(&sigBytes); err != nil {
			return nil, err
		}

		m[txid] = hex.EncodeToString(sigBytes.Bytes())
	}
	return m, nil
}

func (s TreePartialSigs) MarshalJSON() ([]byte, error) {
	mapObject := make(map[string]string)
	for txid, sig := range s {
		var sigBytes bytes.Buffer
		if err := sig.Encode(&sigBytes); err != nil {
			return nil, err
		}

		mapObject[txid] = hex.EncodeToString(sigBytes.Bytes())
	}

	return json.Marshal(mapObject)
}

func (s *TreePartialSigs) UnmarshalJSON(data []byte) error {
	mapObject := make(map[string]string)
	if err := json.Unmarshal(data, &mapObject); err != nil {
		return err
	}

	for txid, sig := range mapObject {
		sigBytes, err := hex.DecodeString(sig)
		if err != nil {
			return err
		}

		sig := &musig2.PartialSignature{}
		if err := sig.Decode(bytes.NewReader(sigBytes)); err != nil {
			return err
		}

		(*s)[txid] = sig
	}
	return nil
}

type SignerSession interface {
	Init(batchOutSweepClosure []byte, batchOutAmount int64, vtxoTree *TxTree) error
	GetPublicKey() string
	GetNonces() (TreeNonces, error) // generate tree nonces for this session
	// deprecated, safer to use AggregateNonces instead
	SetAggregatedNonces(TreeNonces) // set the aggregated nonces
	AggregateNonces(txid string, pubkeyNonces map[string]*Musig2Nonce) (hasAllNonces bool, err error)
	Sign() (TreePartialSigs, error) // sign the tree
}

type CoordinatorSession interface {
	AddNonce(*btcec.PublicKey, TreeNonces)
	AddSignatures(*btcec.PublicKey, TreePartialSigs) (shouldBan bool, err error)
	AggregateNonces() (TreeNonces, error)
	GetPublicNonces() map[string]TreeNonces // pubkey -> nonces
	// SignTree combines the signatures and add them to the tree's psbts
	SignTree() (*TxTree, error)
}

// AggregateKeys is a wrapper around musig2.AggregateKeys using the given tweak
func AggregateKeys(pubkeys []*btcec.PublicKey, tweak []byte) (*musig2.AggregateKey, error) {
	if len(pubkeys) == 0 {
		return nil, errors.New("no pubkeys")
	}

	for _, pubkey := range pubkeys {
		if pubkey == nil {
			return nil, errors.New("nil pubkey")
		}
	}

	// if there is only one pubkey, fallback to classic P2TR
	if len(pubkeys) == 1 {
		res := &musig2.AggregateKey{
			PreTweakedKey: pubkeys[0],
		}

		if len(tweak) > 0 {
			finalKey := txscript.ComputeTaprootOutputKey(pubkeys[0], tweak)
			res.FinalKey = finalKey
		} else {
			res.FinalKey = pubkeys[0]
		}

		return res, nil
	}

	opts := make([]musig2.KeyAggOption, 0)
	if len(tweak) > 0 {
		opts = append(opts, musig2.WithTaprootKeyTweak(tweak))
	}

	key, _, _, err := musig2.AggregateKeys(pubkeys, true, opts...)
	if err != nil {
		return nil, err
	}

	return key, nil
}

// ValidateTreeSigs validates the signatures on the given vtxo tree.
//
// It validates that each input in the tree has a valid signature, and that the
// signature was generated using the aggregated public key of all cosigners.
//
// The function returns an error if any of the signatures are invalid, or if there
// are any issues with the aggregated public key.
func ValidateTreeSigs(
	batchOutSweepClosure []byte, batchOutAmount int64, vtxoTree *TxTree,
) error {
	prevoutFetcherFactory, err := prevOutFetcherFactory(
		batchOutSweepClosure, batchOutAmount, vtxoTree,
	)
	if err != nil {
		return err
	}

	txs := treeToIndexedTxs(vtxoTree, make(map[string]*psbt.Packet))

	_, err = workPoolMap(txs, func(ptx *psbt.Packet) (any, error) {
		sig := ptx.Inputs[0].TaprootKeySpendSig
		if len(sig) == 0 {
			return nil, errors.New("unsigned tree input")
		}

		schnorrSig, err := schnorr.ParseSignature(sig)
		if err != nil {
			return nil, fmt.Errorf("failed to parse signature: %w", err)
		}

		prevoutFetcher, err := prevoutFetcherFactory(ptx)
		if err != nil {
			return nil, fmt.Errorf("failed to get prevout fetcher: %w", err)
		}

		message, err := txscript.CalcTaprootSignatureHash(
			txscript.NewTxSigHashes(ptx.UnsignedTx, prevoutFetcher),
			txscript.SigHashDefault, ptx.UnsignedTx, 0, prevoutFetcher,
		)
		if err != nil {
			return nil, fmt.Errorf("failed to calculate sighash: %w", err)
		}

		consignerPubkeys, err := txutils.ParseCosignerKeysFromArkPsbt(ptx, 0)
		if err != nil {
			return nil, fmt.Errorf(
				"failed to extract cosigners from tx %s: %s", ptx.UnsignedTx.TxID(), err,
			)
		}
		if len(consignerPubkeys) == 0 {
			return nil, fmt.Errorf("no cosigner pubkeys found in tx %s", ptx.UnsignedTx.TxID())
		}

		aggregateKey, err := AggregateKeys(consignerPubkeys, batchOutSweepClosure)
		if err != nil {
			return nil, fmt.Errorf("failed to aggregate keys: %w", err)
		}

		if !schnorrSig.Verify(message, aggregateKey.FinalKey) {
			return nil, fmt.Errorf("invalid signature for txid %s", ptx.UnsignedTx.TxID())
		}

		return nil, nil
	})
	return err
}

func NewTreeSignerSession(signer *btcec.PrivateKey) SignerSession {
	return &treeSignerSession{secretKey: signer}
}

type treeSignerSession struct {
	secretKey             *btcec.PrivateKey
	txs                   map[string]*psbt.Packet
	myNonces              map[string]*musig2.Nonces
	aggregateNonces       TreeNonces
	scriptRoot            []byte
	prevoutFetcherFactory func(*psbt.Packet) (txscript.PrevOutputFetcher, error)
}

func (t *treeSignerSession) Init(
	batchOutSweepClosure []byte, batchOutAmount int64, vtxoTree *TxTree,
) error {
	prevOutFetcherFactory, err := prevOutFetcherFactory(
		batchOutSweepClosure, batchOutAmount, vtxoTree,
	)
	if err != nil {
		return err
	}

	t.scriptRoot = batchOutSweepClosure
	t.prevoutFetcherFactory = prevOutFetcherFactory
	t.txs = treeToIndexedTxs(vtxoTree, make(map[string]*psbt.Packet))
	return nil
}

func (t *treeSignerSession) GetPublicKey() string {
	return hex.EncodeToString(t.secretKey.PubKey().SerializeCompressed())
}

// GetNonces returns only the public musig2 nonces for each transaction
// where the signer's key is in the list of cosigners
func (t *treeSignerSession) GetNonces() (TreeNonces, error) {
	if t.myNonces == nil {
		if err := t.generateNonces(); err != nil {
			return nil, err
		}
	}

	publicNonces := make(TreeNonces)
	for txid, nonces := range t.myNonces {
		publicNonces[txid] = &Musig2Nonce{nonces.PubNonce}
	}

	return publicNonces, nil
}

func (t *treeSignerSession) SetAggregatedNonces(nonces TreeNonces) {
	t.aggregateNonces = nonces
}

func (t *treeSignerSession) AggregateNonces(txid string, pubkeyNonces map[string]*Musig2Nonce) (bool, error) {
	if t.myNonces == nil {
		return false, fmt.Errorf("nonces not generated")
	}

	if t.txs == nil {
		return false, fmt.Errorf("vtxo tree not initialized")
	}

	tx, ok := t.txs[txid]
	if !ok {
		return false, fmt.Errorf("missing tx %s", txid)
	}

	myNonce, ok := t.myNonces[txid]
	if !ok {
		return false, fmt.Errorf("missing my nonce for txid %s", txid)
	}

	if t.aggregateNonces == nil {
		t.aggregateNonces = make(TreeNonces)
	}

	if aggNonce, ok := t.aggregateNonces[txid]; ok && aggNonce != nil {
		return len(t.aggregateNonces) == len(t.txs), nil // skip: we already have the aggregated nonce for this txid
	}

	// set our nonce to ensure server didn't overwrite it
	xonlyPubkey := hex.EncodeToString(schnorr.SerializePubKey(t.secretKey.PubKey()))
	pubkeyNonces[xonlyPubkey] = &Musig2Nonce{myNonce.PubNonce}

	keys, err := txutils.ParseCosignerKeysFromArkPsbt(tx, 0)
	if err != nil {
		return false, fmt.Errorf("failed to get cosigner keys: %w", err)
	}

	if len(keys) == 0 {
		return false, fmt.Errorf("no keys for txid %s", tx.UnsignedTx.TxID())
	}

	allNonces := make([][66]byte, 0, len(keys))
	for _, key := range keys {
		keyStr := hex.EncodeToString(schnorr.SerializePubKey(key))
		nonce, ok := pubkeyNonces[keyStr]
		if !ok {
			return false, fmt.Errorf("missing nonce for cosigner key %s", keyStr)
		}
		allNonces = append(allNonces, nonce.PubNonce)
	}

	aggregatedNonce, err := musig2.AggregateNonces(allNonces)
	if err != nil {
		return false, fmt.Errorf("failed to aggregate nonces: %w", err)
	}

	t.aggregateNonces[txid] = &Musig2Nonce{aggregatedNonce}
	return len(t.aggregateNonces) == len(t.txs), nil
}

// Sign generates the musig2 partial signatures for each transaction of the tree that includes the
// signer as in the cosigners list.
func (t *treeSignerSession) Sign() (TreePartialSigs, error) {
	if t.txs == nil {
		return nil, errors.New("vtxo tree not initialized")
	}

	if t.aggregateNonces == nil {
		return nil, errors.New("nonces not set")
	}

	musigParamsMap := make(map[string]musigParams)
	serializedSignerPubKey := schnorr.SerializePubKey(t.secretKey.PubKey())

	for txid, tx := range t.txs {
		mustSign, cosigners, err := getCosignersPublicKeys(serializedSignerPubKey, tx)
		if err != nil {
			return nil, err
		}

		// if the signer's key is not in the list of keys, skip signing
		if !mustSign {
			continue
		}

		prevoutFetcher, err := t.prevoutFetcherFactory(tx)
		if err != nil {
			return nil, err
		}

		combinedNonce, ok := t.aggregateNonces[txid]
		if !ok {
			return nil, fmt.Errorf("missing combined nonce for txid %s", txid)
		}

		secretNonce, ok := t.myNonces[txid]
		if !ok {
			return nil, fmt.Errorf("missing secret nonce for txid %s", txid)
		}

		musigParamsMap[txid] = musigParams{
			tx:             tx,
			prevoutFetcher: prevoutFetcher,
			combinedNonce:  combinedNonce.PubNonce,
			secretNonce:    secretNonce.SecNonce,
			cosigners:      cosigners,
		}
	}

	return workPoolMap(musigParamsMap, sign(t.secretKey, t.scriptRoot))
}

// generateNonces generates musig2 private and public nonces for each transaction of the tree.
func (t *treeSignerSession) generateNonces() error {
	if len(t.txs) == 0 {
		return ErrMissingVtxoTree
	}

	myNonces, err := workPoolMap(t.txs, generateNonces(t.secretKey.PubKey()))
	if err != nil {
		return err
	}

	t.myNonces = myNonces
	return nil
}

type treeCoordinatorSession struct {
	scriptRoot            []byte
	nonces                map[string]TreeNonces      // xonly pubkey -> nonces
	sigs                  map[string]TreePartialSigs // xonly pubkey -> sigs
	prevoutFetcherFactory func(*psbt.Packet) (txscript.PrevOutputFetcher, error)
	vtxoTree              *TxTree
	txs                   map[string]*psbt.Packet
	combinedNonces        TreeNonces
}

func NewTreeCoordinatorSession(
	batchOutSweepClosure []byte, batchOutAmount int64, vtxoTree *TxTree,
) (CoordinatorSession, error) {
	prevoutFetcherFactory, err := prevOutFetcherFactory(
		batchOutSweepClosure, batchOutAmount, vtxoTree,
	)
	if err != nil {
		return nil, err
	}

	return &treeCoordinatorSession{
		scriptRoot:            batchOutSweepClosure,
		nonces:                make(map[string]TreeNonces),
		sigs:                  make(map[string]TreePartialSigs),
		prevoutFetcherFactory: prevoutFetcherFactory,
		vtxoTree:              vtxoTree,
		txs:                   treeToIndexedTxs(vtxoTree, make(map[string]*psbt.Packet)),
	}, nil
}

func (t *treeCoordinatorSession) AddNonce(pubkey *btcec.PublicKey, nonce TreeNonces) {
	t.nonces[hex.EncodeToString(schnorr.SerializePubKey(pubkey))] = nonce
}

func (t *treeCoordinatorSession) GetPublicNonces() map[string]TreeNonces {
	return t.nonces
}

func (t *treeCoordinatorSession) AddSignatures(pubkey *btcec.PublicKey, sig TreePartialSigs) (shouldBan bool, err error) {
	nonces, ok := t.nonces[hex.EncodeToString(schnorr.SerializePubKey(pubkey))]
	if !ok {
		return false, fmt.Errorf("missing musig2 nonces for pubkey %s", hex.EncodeToString(schnorr.SerializePubKey(pubkey)))
	}

	for txid, tx := range t.txs {
		serializedPubkey := schnorr.SerializePubKey(pubkey)
		mustSign, cosigners, err := getCosignersPublicKeys(serializedPubkey, tx)
		if err != nil {
			return false, fmt.Errorf("failed to get cosigners public keys: %w", err)
		}

		if !mustSign {
			// the signer doesn't have to sign this tx, skip it
			continue
		}

		sig, ok := sig[txid]
		if !ok {
			return true, fmt.Errorf("missing musig2 signature for tx %s", txid)
		}

		nonce, ok := nonces[txid]
		if !ok {
			return true, fmt.Errorf("missing musig2 nonce for txid %s", txid)
		}

		combinedNonce, ok := t.combinedNonces[txid]
		if !ok {
			return false, fmt.Errorf("missing combined nonce for txid %s, cannot validate signature", txid)
		}

		prevoutFetcher, err := t.prevoutFetcherFactory(tx)
		if err != nil {
			return false, fmt.Errorf("failed to get prevout fetcher: %w", err)
		}

		message, err := txscript.CalcTaprootSignatureHash(
			txscript.NewTxSigHashes(tx.UnsignedTx, prevoutFetcher),
			txscript.SigHashDefault, tx.UnsignedTx, 0, prevoutFetcher,
		)
		if err != nil {
			return false, fmt.Errorf("failed to calculate sighash: %w", err)
		}
		if len(message) != 32 {
			return false, fmt.Errorf("invalid taproot signature hash length for txid %s", tx.UnsignedTx.TxID())
		}

		if !sig.Verify(nonce.PubNonce, combinedNonce.PubNonce, cosigners, pubkey, [32]byte(message), musig2.WithTaprootSignTweak(t.scriptRoot)) {
			return true, fmt.Errorf("invalid signature for txid %s", txid)
		}
	}

	t.sigs[hex.EncodeToString(schnorr.SerializePubKey(pubkey))] = sig
	return false, nil
}

// AggregateNonces aggregates the musig2 nonces for each transaction in the tree
// it returns an error if any of the nonces are not set
func (t *treeCoordinatorSession) AggregateNonces() (TreeNonces, error) {
	combinedNonces, err := workPoolMap(t.txs, combineNonces(t.nonces))
	if err != nil {
		return nil, err
	}

	t.combinedNonces = combinedNonces

	return t.combinedNonces, nil
}

// SignTree combines the signatures and add them to the tree's psbts
// it returns the vtxo tree with the signed transactions set as TaprootKeySpendSig
func (t *treeCoordinatorSession) SignTree() (*TxTree, error) {
	combineSigsItems := make(map[string]combineSigsParams)

	for txid, tx := range t.txs {
		prevoutFetcher, err := t.prevoutFetcherFactory(tx)
		if err != nil {
			return nil, err
		}

		combineSigsItems[txid] = combineSigsParams{tx: tx, prevoutFetcher: prevoutFetcher}
	}

	combinedSigs, err := workPoolMap(combineSigsItems, combineSigs(t.scriptRoot, t.sigs))
	if err != nil {
		return nil, err
	}

	if err := t.vtxoTree.Apply(func(g *TxTree) (bool, error) {
		sig, ok := combinedSigs[g.Root.UnsignedTx.TxID()]
		if !ok {
			return true, nil
		}

		g.Root.Inputs[0].TaprootKeySpendSig = sig.Serialize()
		return true, nil
	}); err != nil {
		return nil, err
	}

	return t.vtxoTree, nil
}

func prevOutFetcherFactory(
	batchOutSweepClosure []byte, batchOutAmount int64, vtxoTree *TxTree,
) (func(*psbt.Packet) (txscript.PrevOutputFetcher, error), error) {
	return func(ptx *psbt.Packet) (txscript.PrevOutputFetcher, error) {
		parentOutpoint := ptx.UnsignedTx.TxIn[0].PreviousOutPoint
		parentTxID := parentOutpoint.Hash.String()
		// root tx case
		if vtxoTree.Root.UnsignedTx.TxIn[0].PreviousOutPoint.Hash.String() == parentTxID {
			cosignerPubkeys, err := txutils.ParseCosignerKeysFromArkPsbt(ptx, 0)
			if err != nil {
				return nil, fmt.Errorf(
					"failed to extract cosigners from tx %s: %s", ptx.UnsignedTx.TxID(), err,
				)
			}
			if len(cosignerPubkeys) == 0 {
				return nil, fmt.Errorf("no cosigner pubkeys found in tx %s", ptx.UnsignedTx.TxID())
			}

			aggregateKey, err := AggregateKeys(cosignerPubkeys, batchOutSweepClosure)
			if err != nil {
				return nil, err
			}

			pkscript, err := script.P2TRScript(aggregateKey.FinalKey)
			if err != nil {
				return nil, err
			}

			return &treePrevOutFetcher{
				prevout: &wire.TxOut{
					Value:    batchOutAmount,
					PkScript: pkscript,
				},
			}, nil
		}

		parent := vtxoTree.Find(parentTxID)

		if parent == nil {
			return nil, errors.New("parent tx not found " + parentTxID)
		}

		return &treePrevOutFetcher{
			prevout: parent.Root.UnsignedTx.TxOut[parentOutpoint.Index],
		}, nil
	}, nil
}

type treePrevOutFetcher struct {
	prevout *wire.TxOut
}

func (f *treePrevOutFetcher) FetchPrevOutput(wire.OutPoint) *wire.TxOut {
	return f.prevout
}

// workPool is a generic worker pool that processes items concurrently
func workPool[T any](items []T, workers int, processItem func(item T) error) error {
	errChan := make(chan error, 1)
	workChan := make(chan T)

	var wg sync.WaitGroup
	wg.Add(workers)

	// launch workers
	for w := 0; w < workers; w++ {
		go func() {
			defer wg.Done()
			for item := range workChan {
				if err := processItem(item); err != nil {
					select {
					case errChan <- err:
					default:
					}
					return
				}
			}
		}()
	}

	// distribute tasks
	go func() {
		for _, item := range items {
			select {
			case err := <-errChan:
				close(workChan)
				errChan <- err
				return
			default:
				workChan <- item
			}
		}
		close(workChan)
	}()

	// wait for all workers to finish
	wg.Wait()

	select {
	case err := <-errChan:
		return err
	default:
		return nil
	}
}

// workPoolMap is a specialized version of workPool for processing maps jobs in parallel
func workPoolMap[T any, R comparable](
	kvMap map[string]T, processItem func(item T) (R, error),
) (map[string]R, error) {
	locker := sync.Mutex{}
	results := make(map[string]R)

	type workItem struct {
		key  string
		item T
	}

	items := make([]workItem, 0, len(kvMap))
	for key, item := range kvMap {
		items = append(items, workItem{key: key, item: item})
	}

	if err := workPool(items, runtime.NumCPU(), func(item workItem) error {
		result, err := processItem(item.item)
		if err != nil {
			return err
		}

		zero := new(R)
		if result != *zero {
			locker.Lock()
			results[item.key] = result
			locker.Unlock()
		}
		return nil
	}); err != nil {
		return nil, err
	}

	return results, nil
}

// getCosignersPublicKeys checks if the given signer public key is present in the cosigner keys for
// the transaction's first input.
// It returns a boolean indicating the presence of the key, the list of cosigner public keys, and
// an error if one occurs while fetching cosigner keys.
func getCosignersPublicKeys(
	signerPubkey []byte, tx *psbt.Packet,
) (bool, []*btcec.PublicKey, error) {
	cosignerFields, err := txutils.GetArkPsbtFields(tx, 0, txutils.CosignerPublicKeyField)
	if err != nil {
		return false, nil, err
	}

	for _, field := range cosignerFields {
		if bytes.Equal(schnorr.SerializePubKey(field.PublicKey), signerPubkey) {
			return true, txutils.ParseCosignersToECPubKeys(cosignerFields), nil
		}
	}
	return false, nil, nil
}

func treeToIndexedTxs(txTree *TxTree, res map[string]*psbt.Packet) map[string]*psbt.Packet {
	res[txTree.Root.UnsignedTx.TxID()] = txTree.Root

	for _, child := range txTree.Children {
		res = treeToIndexedTxs(child, res)
	}

	return res
}

func combineNonces(allNonces map[string]TreeNonces) func(tx *psbt.Packet) (*Musig2Nonce, error) {
	return func(tx *psbt.Packet) (*Musig2Nonce, error) {
		keys, err := txutils.GetArkPsbtFields(tx, 0, txutils.CosignerPublicKeyField)
		if err != nil {
			return nil, fmt.Errorf("failed to get cosigner keys: %w", err)
		}

		if len(keys) == 0 {
			return nil, fmt.Errorf("no keys for txid %s", tx.UnsignedTx.TxID())
		}

		nonces := make([][66]byte, 0, len(keys))

		for _, field := range keys {
			keyStr := hex.EncodeToString(schnorr.SerializePubKey(field.PublicKey))
			nonceMap, ok := allNonces[keyStr]
			if !ok {
				return nil, fmt.Errorf(
					"nonces not set for cosigner key %x",
					field.PublicKey.SerializeCompressed(),
				)
			}

			nonce := nonceMap[tx.UnsignedTx.TxID()]
			if nonce == nil {
				return nil, fmt.Errorf(
					"missing nonce for cosigner key %x",
					field.PublicKey.SerializeCompressed(),
				)
			}

			nonces = append(nonces, nonce.PubNonce)
		}

		if len(nonces) == 0 {
			return nil, fmt.Errorf("missing nonces for txid %s", tx.UnsignedTx.TxID())
		}

		if len(nonces) != len(keys) {
			return nil, fmt.Errorf(
				"wrong number of nonces for txid %s, expected %d got %d",
				tx.UnsignedTx.TxID(),
				len(keys),
				len(nonces),
			)
		}

		aggregatedNonce, err := musig2.AggregateNonces(nonces)
		if err != nil {
			return nil, fmt.Errorf("failed to aggregate nonces: %w", err)
		}

		return &Musig2Nonce{aggregatedNonce}, nil
	}
}

func generateNonces(
	signerPubKey *btcec.PublicKey,
) func(*psbt.Packet) (*musig2.Nonces, error) {
	serializedSignerPubKey := schnorr.SerializePubKey(signerPubKey)

	return func(ptx *psbt.Packet) (*musig2.Nonces, error) {
		mustGenerateNonce, _, err := getCosignersPublicKeys(serializedSignerPubKey, ptx)
		if err != nil {
			return nil, err
		}

		// if the signer's key is not in the list of keys, skip generating nonces
		if !mustGenerateNonce {
			return nil, nil
		}

		// generate musig2 nonces
		nonce, err := musig2.GenNonces(
			musig2.WithPublicKey(signerPubKey),
		)
		if err != nil {
			return nil, err
		}

		return nonce, nil
	}
}

type musigParams struct {
	tx             *psbt.Packet
	combinedNonce  [66]byte
	secretNonce    [97]byte
	prevoutFetcher txscript.PrevOutputFetcher
	cosigners      []*btcec.PublicKey
}

func sign(
	signer *btcec.PrivateKey, batchOutSweepClosure []byte,
) func(musigParams) (*musig2.PartialSignature, error) {
	return func(params musigParams) (*musig2.PartialSignature, error) {
		message, err := txscript.CalcTaprootSignatureHash(
			txscript.NewTxSigHashes(params.tx.UnsignedTx, params.prevoutFetcher),
			txscript.SigHashDefault, params.tx.UnsignedTx, 0, params.prevoutFetcher,
		)
		if err != nil {
			return nil, err
		}
		if len(message) != 32 {
			return nil, fmt.Errorf("invalid taproot signature hash length for txid %s", params.tx.UnsignedTx.TxID())
		}

		return musig2.Sign(
			params.secretNonce, signer, params.combinedNonce, params.cosigners, [32]byte(message),
			musig2.WithSortedKeys(), musig2.WithTaprootSignTweak(batchOutSweepClosure),
			musig2.WithFastSign(),
		)
	}
}

type combineSigsParams struct {
	tx             *psbt.Packet
	prevoutFetcher txscript.PrevOutputFetcher
}

// combineSigs combines the partial signatures from all cosigners
// and returns the combined schnorr signature.
//
// The function also takes the prevout fetcher, which is used to fetch
// the prevout for the given txid.
//
// The function returns an error if the partial signatures are not
// set for any of the cosigners, or if the combined nonce is not set.
//
// The function returns an error if the combined nonce is not set, or
// if the signatures are invalid.
//
// The function returns an error if the number of signatures is not
// equal to the number of cosigners.
//
// The function returns an error if the aggregated key is not set.
func combineSigs(
	batchOutSweepClosure []byte, allSigs map[string]TreePartialSigs,
) func(combineSigsParams) (*schnorr.Signature, error) {
	return func(params combineSigsParams) (*schnorr.Signature, error) {
		keys, err := txutils.ParseCosignerKeysFromArkPsbt(params.tx, 0)
		if err != nil {
			return nil, fmt.Errorf(
				"failed to extract cosigners from tx %s: %s", params.tx.UnsignedTx.TxID(), err,
			)
		}
		if len(keys) == 0 {
			return nil, fmt.Errorf(
				"no cosigner pubkeys found in tx %s", params.tx.UnsignedTx.TxID(),
			)
		}

		var combinedNonce *btcec.PublicKey
		sigs := make([]*musig2.PartialSignature, 0, len(keys))

		for _, key := range keys {
			keySigs, ok := allSigs[hex.EncodeToString(schnorr.SerializePubKey(key))]
			if !ok {
				return nil, fmt.Errorf(
					"sigs not set for cosigner key %x", key.SerializeCompressed(),
				)
			}

			s := keySigs[params.tx.UnsignedTx.TxID()]
			if s == nil {
				return nil, fmt.Errorf(
					"missing signature for cosigner key %x", key.SerializeCompressed(),
				)
			}

			if s.R != nil {
				combinedNonce = s.R
			}
			sigs = append(sigs, s)
		}

		if combinedNonce == nil {
			return nil, fmt.Errorf(
				"missing combined nonce for txid %s",
				params.tx.UnsignedTx.TxID(),
			)
		}

		message, err := txscript.CalcTaprootSignatureHash(
			txscript.NewTxSigHashes(params.tx.UnsignedTx, params.prevoutFetcher),
			txscript.SigHashDefault, params.tx.UnsignedTx, 0, params.prevoutFetcher,
		)
		if err != nil {
			return nil, err
		}
		if len(message) != 32 {
			return nil, fmt.Errorf("invalid taproot signature hash length for txid %s", params.tx.UnsignedTx.TxID())
		}

		if len(sigs) == 0 {
			return nil, fmt.Errorf("missing signatures for txid %s", params.tx.UnsignedTx.TxID())
		}

		if len(sigs) != len(keys) {
			return nil, fmt.Errorf(
				"wrong number of signatures for txid %s, expected %d got %d",
				params.tx.UnsignedTx.TxID(), len(keys), len(sigs),
			)
		}

		combineOpts := make([]musig2.CombineOption, 0)
		if batchOutSweepClosure != nil {
			combineOpts = append(
				combineOpts, musig2.WithTaprootTweakedCombine(
					[32]byte(message), keys, batchOutSweepClosure, true,
				),
			)
		}

		combinedSig := musig2.CombineSigs(
			combinedNonce, sigs,
			combineOpts...,
		)

		return combinedSig, nil
	}
}
