package experimental_test

import (
	"context"
	_ "embed"
	"fmt"
	"log"

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

// snapshotWasm was generated by the following:
//
//	cd testdata; wat2wasm snapshot.wat
//
//go:embed testdata/snapshot.wasm
var snapshotWasm []byte

type snapshotsKey struct{}

func Example_enableSnapshotterKey() {
	ctx := context.Background()

	rt := wazero.NewRuntime(ctx)
	defer rt.Close(ctx) // This closes everything this Runtime created.

	// Enable experimental snapshotting functionality by setting it to context. We use this
	// context when invoking functions, indicating to wazero to enable it.
	ctx = experimental.WithSnapshotter(ctx)

	// Also place a mutable holder of snapshots to be referenced during restore.
	var snapshots []experimental.Snapshot
	ctx = context.WithValue(ctx, snapshotsKey{}, &snapshots)

	// Register host functions using snapshot and restore. Generally snapshot is saved
	// into a mutable location in context to be referenced during restore.
	_, err := rt.NewHostModuleBuilder("example").
		NewFunctionBuilder().
		WithFunc(func(ctx context.Context, mod api.Module, snapshotPtr uint32) int32 {
			// Because we enabled snapshots with WithSnapshotter, this is non-nil.
			snapshot := experimental.GetSnapshotter(ctx).Snapshot()

			// Get our mutable snapshots holder to be able to add to it. Our example only calls snapshot
			// and restore once but real programs will often call them at multiple layers within a call
			// stack with various e.g., try/catch statements.
			snapshots := ctx.Value(snapshotsKey{}).(*[]experimental.Snapshot)
			idx := len(*snapshots)
			*snapshots = append(*snapshots, snapshot)

			// Write a value to be passed back to restore. This is meant to be opaque to the guest
			// and used to re-reference the snapshot.
			ok := mod.Memory().WriteUint32Le(snapshotPtr, uint32(idx))
			if !ok {
				log.Panicln("failed to write snapshot index")
			}

			return 0
		}).
		Export("snapshot").
		NewFunctionBuilder().
		WithFunc(func(ctx context.Context, mod api.Module, snapshotPtr uint32) {
			// Read the value written by snapshot to re-reference the snapshot.
			idx, ok := mod.Memory().ReadUint32Le(snapshotPtr)
			if !ok {
				log.Panicln("failed to read snapshot index")
			}

			// Get the snapshot
			snapshots := ctx.Value(snapshotsKey{}).(*[]experimental.Snapshot)
			snapshot := (*snapshots)[idx]

			// Restore! The invocation of this function will end as soon as we invoke
			// Restore, so we also pass in our return value. The guest function run
			// will finish with this return value.
			snapshot.Restore([]uint64{5})
		}).
		Export("restore").
		Instantiate(ctx)
	if err != nil {
		log.Panicln(err)
	}

	mod, err := rt.Instantiate(ctx, snapshotWasm) // Instantiate the actual code
	if err != nil {
		log.Panicln(err)
	}

	// Call the guest entrypoint.
	res, err := mod.ExportedFunction("run").Call(ctx)
	if err != nil {
		log.Panicln(err)
	}
	// We restored and returned the restore value, so it's our result. If restore
	// was instead a no-op, we would have returned 10 from normal code flow.
	fmt.Println(res[0])
	// Output:
	// 5
}