package nonce

import (
	"context"
	"crypto/ecdsa"
	"fmt"
	"github.com/synapsecns/sanguine/ethergo/util"
	"math/big"
	"sync"

	"github.com/ethereum/go-ethereum/accounts/abi/bind"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/synapsecns/sanguine/core"
	"github.com/synapsecns/sanguine/core/mapmutex"
)

// Manager is a singleton used to generate a nonce.
// This can be explicitly overrided by setting a nonce manually (not recommended). The issue here is many of the
// deploy helper functions are called asynchronously and there's a gap between the tx.Sequence generation and the
// transactor. This solves that by wrapping the transactor in a nonce manager.
type Manager interface {
	// SignTx signs a legacy tx
	SignTx(ogTx *types.Transaction, signer types.Signer, prv *ecdsa.PrivateKey, options ...Option) (*types.Transaction, error)
	// NewKeyedTransactor wraps keyed transactor in a nonce manager.
	// right now, this only works if all txes are sent out (a safe assumption in test mode)
	// this can be obviated by signing at send time or loop + retrying on failure
	NewKeyedTransactor(realSigner *bind.TransactOpts) (*bind.TransactOpts, error)
	// NewKeyedTransactorFromKey creates a new keyed transactor from a private key.
	NewKeyedTransactorFromKey(key *ecdsa.PrivateKey) (*bind.TransactOpts, error)
	// GetNextNonce gets the next nonce for the address.
	GetNextNonce(address common.Address) (*big.Int, error)
	// ClearNonce clears the nonce for the address.
	ClearNonce(address common.Address)
}

// ChainQuery is a chain used to generate a nonce.
type ChainQuery interface {
	// PendingNonceAt returns the account nonce of the given account in the pending state.
	// This is the nonce that should be used for the next transaction.
	PendingNonceAt(ctx context.Context, account common.Address) (uint64, error)
}

type nonceManagerImp struct {
	// chain is the chain to use for the nonce
	chain ChainQuery
	// ctx is the context used while fetching data
	//nolint: containedctx
	ctx context.Context
	// chainID is the chain id of the nonce manager.
	chainID *big.Int
	// nonceMapLock locks the nonce map for reads/writes. We use this instead of a sync
	// map because of type safety
	nonceMapLock sync.RWMutex
	// nonceMap is a map of accounts->next nonce
	nonceMap map[common.Address]*big.Int
	// accountMutex is responsible for making sure only one nonce per account is generated at a time
	// we don't use the global nonce lock here since we want to allow others to sign transactions during
	// signing. We can't paralellize this per account because the transactor may need to subtract if signing fails
	accountMutex mapmutex.StringerMapMutex
}

// NewNonceManager generates a new nonce manager. This should be called once, as nonce storage is not global.
func NewNonceManager(ctx context.Context, chain ChainQuery, chainID *big.Int) Manager {
	return &nonceManagerImp{
		//nolint: containedctx
		ctx:          ctx,
		chain:        chain,
		chainID:      core.CopyBigInt(chainID),
		nonceMap:     make(map[common.Address]*big.Int),
		nonceMapLock: sync.RWMutex{},
		accountMutex: mapmutex.NewStringerMapMutex(),
	}
}

// GetNextNonce gets the next nonce for the account.
func (n *nonceManagerImp) GetNextNonce(address common.Address) (*big.Int, error) {
	// get the next nonce for the account.
	n.nonceMapLock.Lock()
	currentNonce := n.nonceMap[address]
	defer n.nonceMapLock.Unlock()

	if currentNonce == nil {
		pendingNonce, err := n.chain.PendingNonceAt(n.ctx, address)
		if err != nil {
			return nil, fmt.Errorf("could not get pending nonce for address %s: %w", address.String(), err)
		}

		logger.Debugf("got first nonce %d for account %s", pendingNonce, address)

		// set the nonce in the map
		currentNonce = big.NewInt(int64(pendingNonce))

		n.nonceMap[address] = currentNonce
	}

	return currentNonce, nil
}

// ClearNonce clears the nonce for the account.
func (n *nonceManagerImp) ClearNonce(address common.Address) {
	// clear the nonce for the account
	n.nonceMapLock.Lock()
	n.nonceMap[address] = nil
	defer n.nonceMapLock.Unlock()
}

// incrementNonce increments the nonce for an account. This should be called from within a accountMutex.
func (n *nonceManagerImp) incrementNonce(address common.Address) error {
	currentNonce, err := n.GetNextNonce(address)
	if err != nil {
		return fmt.Errorf("could not get current nonce: %w", err)
	}

	n.nonceMapLock.Lock()
	n.nonceMap[address] = big.NewInt(0).Add(currentNonce, big.NewInt(1))
	n.nonceMapLock.Unlock()

	return nil
}

// NewKeyedTransactorFromKey wraps keyed transactor in a nonce manager.
// right now, this only works if all txes are sent out (a safe assumption in test mode)
// this can be obviated by signing at send time or loop + retrying on failure.
func (n *nonceManagerImp) NewKeyedTransactorFromKey(key *ecdsa.PrivateKey) (*bind.TransactOpts, error) {
	realSigner, err := bind.NewKeyedTransactorWithChainID(key, n.chainID)
	if err != nil {
		return nil, fmt.Errorf("could not create signer: %w", err)
	}
	return n.NewKeyedTransactor(realSigner)
}

// NewKeyedTransactor wraps keyed transactor in a nonce manager.
// right now, this only works if all txes are sent out (a safe assumption in test mode)
// this can be obviated by signing at send time or loop + retrying on failure.
func (n *nonceManagerImp) NewKeyedTransactor(realSigner *bind.TransactOpts) (*bind.TransactOpts, error) {
	return &bind.TransactOpts{
		From: realSigner.From,
		Signer: func(address common.Address, transaction *types.Transaction) (*types.Transaction, error) {
			// lock the account
			acctMutex := n.accountMutex.Lock(address)
			defer acctMutex.Unlock()

			// get the next nonce for the account
			nonce, err := n.GetNextNonce(realSigner.From)
			if err != nil {
				return nil, fmt.Errorf("could not get next nonce: %w", err)
			}

			copiedTx, err := util.CopyTX(transaction, util.WithNonce(nonce.Uint64()))
			if err != nil {
				return nil, fmt.Errorf("could not copy tx: %w", err)
			}
			signedTx, err := realSigner.Signer(address, copiedTx)
			// sign was unsuccesfull, don't increment nonce
			if err != nil {
				return nil, fmt.Errorf("could not sign tx: %w", err)
			}

			err = n.incrementNonce(address)
			if err != nil {
				return nil, err
			}
			return signedTx, nil
		},
	}, nil
}

// SignTx signs a legacy tx.
func (n *nonceManagerImp) SignTx(ogTx *types.Transaction, signer types.Signer, prv *ecdsa.PrivateKey, options ...Option) (*types.Transaction, error) {
	cfg := &signTXConfig{}

	for _, opt := range options {
		opt(cfg)
	}

	address := crypto.PubkeyToAddress(prv.PublicKey)

	addressLock := n.accountMutex.Lock(address)
	defer addressLock.Unlock()

	var copyOpts []util.CopyOption
	if !cfg.skipNonceBump {
		nonce, err := n.GetNextNonce(address)
		if err != nil {
			return nil, fmt.Errorf("could not get nonce: %w", err)
		}

		copyOpts = append(copyOpts, util.WithNonce(nonce.Uint64()))
	}

	tx, err := util.CopyTX(ogTx, copyOpts...)
	if err != nil {
		return nil, fmt.Errorf("could not copy tx: %w", err)
	}

	tx, err = types.SignTx(tx, signer, prv)
	if err != nil {
		return nil, fmt.Errorf("could not sign tx: %w", err)
	}

	if !cfg.skipNonceBump {
		err = n.incrementNonce(address)
		if err != nil {
			return nil, fmt.Errorf("could not increment nonce: %w", err)
		}
	}

	return tx, nil
}

// Option is a functional option for SignTX method.
type Option func(*signTXConfig)

// signTXConfig is the config for SignTX.
type signTXConfig struct {
	skipNonceBump bool
}

// WithNoBump sets the skipNonceBump flag.
func WithNoBump(noBump bool) Option {
	return func(config *signTXConfig) {
		config.skipNonceBump = noBump
	}
}