// Package vm provides a WASM-based smart contract execution engine.
//
// The engine uses wazero (pure-Go, no CGO) in interpreter mode to guarantee
// deterministic execution across all platforms — a requirement for consensus.
//
// Contract lifecycle:
//  1. DEPLOY_CONTRACT tx → chain calls Validate() to compile-check the WASM,
//     then stores the bytecode in BadgerDB.
//  2. CALL_CONTRACT tx   → chain calls Call() with the stored WASM bytes,
//     the method name, JSON args, and a gas limit.
//
// Gas model: each WASM function call costs gasPerCall (100) units.
// Gas cost in µT = gasUsed × blockchain.GasPrice.
package vm

import (
	"context"
	"fmt"
	"log"
	"sync"

	"github.com/tetratelabs/wazero"
	"github.com/tetratelabs/wazero/experimental"
	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"

	"go-blockchain/blockchain"
)

// VM is a WASM execution engine. Create one per process; it is safe for
// concurrent use.  Compiled modules are cached by contract ID.
type VM struct {
	rt    wazero.Runtime
	mu    sync.RWMutex
	cache map[string]wazero.CompiledModule // contractID → compiled module
}

// NewVM creates a VM with an interpreter-mode wazero runtime.
// Interpreter mode guarantees identical behaviour on all OS/arch combos,
// which is required for all nodes to reach the same state.
//
// WASI preview 1 is pre-instantiated to support contracts compiled with
// TinyGo's wasip1 target (tinygo build -target wasip1). Contracts that do
// not import from wasi_snapshot_preview1 are unaffected.
func NewVM(ctx context.Context) *VM {
	cfg := wazero.NewRuntimeConfigInterpreter()
	cfg = cfg.WithDebugInfoEnabled(false)
	// Enable cooperative cancellation: fn.Call(ctx) will return when ctx is
	// cancelled, even if the contract is in an infinite loop. Without this,
	// a buggy or malicious contract that dodges gas metering (e.g. a tight
	// loop of opcodes not hooked by the gas listener) would hang the
	// AddBlock goroutine forever, freezing the entire chain.
	cfg = cfg.WithCloseOnContextDone(true)
	// Attach call-level gas metering.
	ctx = experimental.WithFunctionListenerFactory(ctx, gasListenerFactory{})
	rt := wazero.NewRuntimeWithConfig(ctx, cfg)
	// Instantiate WASI so TinyGo contracts can call proc_exit and basic I/O.
	// The sandbox is fully isolated — no filesystem or network access.
	wasi_snapshot_preview1.MustInstantiate(ctx, rt)
	return &VM{
		rt:    rt,
		cache: make(map[string]wazero.CompiledModule),
	}
}

// Close releases the underlying wazero runtime.
func (v *VM) Close(ctx context.Context) error {
	return v.rt.Close(ctx)
}

// Validate compiles the WASM bytes without executing them.
// Returns an error if the bytes are not valid WASM or import unknown symbols.
// Implements blockchain.ContractVM.
func (v *VM) Validate(ctx context.Context, wasmBytes []byte) error {
	_, err := v.rt.CompileModule(ctx, wasmBytes)
	if err != nil {
		return fmt.Errorf("invalid WASM module: %w", err)
	}
	return nil
}

// Call executes method on the contract identified by contractID.
// wasmBytes is the compiled WASM (loaded from DB by the chain).
// Returns gas consumed. Returns ErrOutOfGas (wrapping ErrTxFailed) on exhaustion.
// Implements blockchain.ContractVM.
func (v *VM) Call(
	ctx context.Context,
	contractID string,
	wasmBytes []byte,
	method string,
	argsJSON []byte,
	gasLimit uint64,
	env blockchain.VMHostEnv,
) (gasUsed uint64, err error) {
	// Attach gas counter to context.
	ctx, gc := withGasCounter(ctx, gasLimit)
	ctx = experimental.WithFunctionListenerFactory(ctx, gasListenerFactory{})

	// Compile (or retrieve from cache).
	compiled, err := v.compiled(ctx, contractID, wasmBytes)
	if err != nil {
		return 0, fmt.Errorf("compile contract %s: %w", contractID, err)
	}

	// Instantiate the "env" host module for this call, wiring in typed args.
	hostInst, err := registerHostModule(ctx, v.rt, env, argsJSON)
	if err != nil {
		return 0, fmt.Errorf("register host module: %w", err)
	}
	defer hostInst.Close(ctx)

	// Instantiate the contract module.
	modCfg := wazero.NewModuleConfig().
		WithName(""). // anonymous — allows multiple concurrent instances
		WithStartFunctions() // do not auto-call _start
	mod, err := v.rt.InstantiateModule(ctx, compiled, modCfg)
	if err != nil {
		return gc.Used(), fmt.Errorf("instantiate contract: %w", err)
	}
	defer mod.Close(ctx)

	// Look up the exported method.
	fn := mod.ExportedFunction(method)
	if fn == nil {
		return gc.Used(), fmt.Errorf("%w: method %q not exported by contract %s",
			blockchain.ErrTxFailed, method, contractID)
	}

	// Call. WASM functions called from Go pass args via the stack; our contracts
	// use no parameters — all input/output goes through host state functions.
	_, callErr := fn.Call(ctx)
	gasUsed = gc.Used()

	if callErr != nil {
		log.Printf("[VM] contract %s.%s error: %v", contractID[:8], method, callErr)
		if isOutOfGas(gc) {
			return gasUsed, ErrOutOfGas
		}
		return gasUsed, fmt.Errorf("%w: %v", blockchain.ErrTxFailed, callErr)
	}
	if isOutOfGas(gc) {
		return gasUsed, ErrOutOfGas
	}
	return gasUsed, nil
}

// compiled returns a cached compiled module, compiling it first if not cached.
// Before compiling, the WASM bytes are instrumented with loop-header gas_tick
// calls so that infinite loops are bounded by the gas limit.
func (v *VM) compiled(ctx context.Context, contractID string, wasmBytes []byte) (wazero.CompiledModule, error) {
	v.mu.RLock()
	cm, ok := v.cache[contractID]
	v.mu.RUnlock()
	if ok {
		return cm, nil
	}

	v.mu.Lock()
	defer v.mu.Unlock()
	// Double-check after acquiring write lock.
	if cm, ok = v.cache[contractID]; ok {
		return cm, nil
	}
	// Instrument: inject gas_tick at loop headers.
	// On instrumentation failure, fall back to the original bytes so that
	// unusual WASM features do not prevent execution entirely.
	instrumented, err := Instrument(wasmBytes)
	if err != nil {
		log.Printf("[VM] instrument contract %s: %v (using original bytes)", contractID[:min8(contractID)], err)
		instrumented = wasmBytes
	}
	compiled, err := v.rt.CompileModule(ctx, instrumented)
	if err != nil {
		return nil, err
	}
	v.cache[contractID] = compiled
	return compiled, nil
}

func min8(s string) int {
	if len(s) < 8 {
		return len(s)
	}
	return 8
}