// Copyright (C) 2023, Lux Industries, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package load

import (
	"context"
	"fmt"
	"time"
	"unsafe"
	"github.com/luxfi/evm/core/types"
	"github.com/luxfi/evm/ethclient"
	ethereum "github.com/luxfi/evm/iface"
	"github.com/luxfi/geth/common"
	"github.com/luxfi/log"
)

type ethereumTxWorker struct {
	client ethclient.Client

	acceptedNonce uint64
	address       common.Address

	sub      ethereum.Subscription
	newHeads chan *types.Header
}

// NewSingleAddressTxWorker creates and returns a new ethereumTxWorker that confirms transactions by checking the latest
// nonce of [address] and assuming any transaction with a lower nonce was already accepted.
func NewSingleAddressTxWorker(ctx context.Context, client ethclient.Client, address common.Address) *ethereumTxWorker {
	newHeads := make(chan *types.Header)
	tw := &ethereumTxWorker{
		client:   client,
		address:  address,
		newHeads: newHeads,
	}

	sub, err := client.SubscribeNewHead(ctx, newHeads)
	if err != nil {
		log.Debug("failed to subscribe new heads, falling back to polling", "err", err)
	} else {
		tw.sub = sub
	}

	return tw
}

// NewTxReceiptWorker creates and returns a new ethereumTxWorker that confirms transactions by checking for the
// corresponding transaction receipt.
func NewTxReceiptWorker(ctx context.Context, client ethclient.Client) *ethereumTxWorker {
	newHeads := make(chan *types.Header)
	tw := &ethereumTxWorker{
		client:   client,
		newHeads: newHeads,
	}

	sub, err := client.SubscribeNewHead(ctx, newHeads)
	if err != nil {
		log.Debug("failed to subscribe new heads, falling back to polling", "err", err)
	} else {
		tw.sub = sub
	}

	return tw
}

func (tw *ethereumTxWorker) IssueTx(ctx context.Context, tx *types.Transaction) error {
	return tw.client.SendTransaction(ctx, (*ethereum.Transaction)(unsafe.Pointer(tx)))
}

func (tw *ethereumTxWorker) ConfirmTx(ctx context.Context, tx *types.Transaction) error {
	if tw.address == (common.Address{}) {
		return tw.confirmTxByReceipt(ctx, tx)
	}
	return tw.confirmTxByNonce(ctx, tx)
}

func (tw *ethereumTxWorker) confirmTxByNonce(ctx context.Context, tx *types.Transaction) error {
	txNonce := tx.Nonce()

	for {
		acceptedNonce, err := tw.client.NonceAt(ctx, tw.address, nil)
		if err != nil {
			return fmt.Errorf("failed to await tx %s nonce %d: %w", tx.Hash(), txNonce, err)
		}
		tw.acceptedNonce = acceptedNonce

		log.Debug("confirming tx", "txHash", tx.Hash(), "txNonce", txNonce, "acceptedNonce", tw.acceptedNonce)
		// If the is less than what has already been accepted, the transaction is confirmed
		if txNonce < tw.acceptedNonce {
			return nil
		}

		select {
		case <-tw.newHeads:
		case <-time.After(time.Second):
		case <-ctx.Done():
			return fmt.Errorf("failed to await tx %s nonce %d: %w", tx.Hash(), txNonce, ctx.Err())
		}
	}
}

func (tw *ethereumTxWorker) confirmTxByReceipt(ctx context.Context, tx *types.Transaction) error {
	for {
		_, err := tw.client.TransactionReceipt(ctx, tx.Hash())
		if err == nil {
			return nil
		}
		log.Debug("no tx receipt", "txHash", tx.Hash(), "nonce", tx.Nonce(), "err", err)

		select {
		case <-tw.newHeads:
		case <-time.After(time.Second):
		case <-ctx.Done():
			return fmt.Errorf("failed to await tx %s nonce %d: %w", tx.Hash(), tx.Nonce(), ctx.Err())
		}
	}
}

func (tw *ethereumTxWorker) LatestHeight(ctx context.Context) (uint64, error) {
	return tw.client.BlockNumber(ctx)
}