package vm

import (
	"context"
	"errors"
	"fmt"
	"sync/atomic"

	"github.com/tetratelabs/wazero/api"
	"github.com/tetratelabs/wazero/experimental"

	"go-blockchain/blockchain"
)

// ErrOutOfGas is returned when a contract call exhausts its gas limit.
// It wraps blockchain.ErrTxFailed so the call is skipped rather than
// aborting the entire block.
var ErrOutOfGas = fmt.Errorf("%w: out of gas", blockchain.ErrTxFailed)

// gasKey is the context key used to pass the gas counter into host functions.
type gasKey struct{}

// gasCounter holds a mutable gas counter accessible through a context value.
type gasCounter struct {
	used  atomic.Uint64
	limit uint64
}

// withGasCounter attaches a new gas counter to ctx.
func withGasCounter(ctx context.Context, limit uint64) (context.Context, *gasCounter) {
	gc := &gasCounter{limit: limit}
	return context.WithValue(ctx, gasKey{}, gc), gc
}

// gasFromContext retrieves the gas counter from ctx; panics if not present.
func gasFromContext(ctx context.Context) *gasCounter {
	gc, _ := ctx.Value(gasKey{}).(*gasCounter)
	return gc
}

// charge attempts to consume n gas units. Returns ErrOutOfGas if the limit is exceeded.
func (gc *gasCounter) charge(n uint64) error {
	if gc == nil {
		return nil
	}
	newUsed := gc.used.Add(n)
	if newUsed > gc.limit {
		return ErrOutOfGas
	}
	return nil
}

// Used returns total gas consumed so far.
func (gc *gasCounter) Used() uint64 {
	if gc == nil {
		return 0
	}
	return gc.used.Load()
}

// Remaining returns gas budget minus gas used. Returns 0 when exhausted.
func (gc *gasCounter) Remaining() uint64 {
	if gc == nil {
		return 0
	}
	used := gc.used.Load()
	if used >= gc.limit {
		return 0
	}
	return gc.limit - used
}

// gasListenerFactory is a wazero FunctionListenerFactory that charges gas on
// every WASM function call. Each function call costs gasPerCall units.
// True instruction-level metering would require bytecode instrumentation;
// call-level metering is a pragmatic approximation that prevents runaway contracts.
const gasPerCall uint64 = 100

type gasListenerFactory struct{}

func (gasListenerFactory) NewFunctionListener(def api.FunctionDefinition) experimental.FunctionListener {
	return gasListener{}
}

type gasListener struct{}

func (gasListener) Before(ctx context.Context, _ api.Module, _ api.FunctionDefinition, _ []uint64, _ experimental.StackIterator) {
	gc := gasFromContext(ctx)
	if gc != nil {
		// Ignore error here — we can't abort from Before.
		// isOutOfGas() is checked after Call() returns.
		_ = gc.charge(gasPerCall)
	}
}

func (gasListener) After(_ context.Context, _ api.Module, _ api.FunctionDefinition, _ []uint64) {}
func (gasListener) Abort(_ context.Context, _ api.Module, _ api.FunctionDefinition, _ error)    {}

// isOutOfGas reports whether err indicates gas exhaustion.
func isOutOfGas(gc *gasCounter) bool {
	if gc == nil {
		return false
	}
	return gc.Used() > gc.limit
}

// ensure gasListenerFactory satisfies the interface at compile time.
var _ experimental.FunctionListenerFactory = gasListenerFactory{}

// ensure errors chain correctly.
var _ = errors.Is(ErrOutOfGas, blockchain.ErrTxFailed)