package test

import (
	rpcv1 "buf.build/gen/go/buildr/buildr/protocolbuffers/go/rpc/v1"
	"context"
	"errors"
	"github.com/tetratelabs/wazero"
	"github.com/tetratelabs/wazero/api"
	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
	"golang.org/x/exp/slog"
	"google.golang.org/protobuf/proto"
)

type Host struct {
	Logger *slog.Logger
}

func (h *Host) Run(ctx context.Context) (err error) {
	r := wazero.NewRuntime(ctx)
	defer r.Close(context.Background())

	_, err = r.NewHostModuleBuilder("test").
		NewFunctionBuilder().WithFunc(h.log).Export("log").
		Instantiate(ctx)

	if err != nil {
		return err
	}

	closer, err := wasi_snapshot_preview1.Instantiate(ctx, r)
	if err != nil {
		return err
	}

	defer func() {
		err = errors.Join(err, closer.Close(context.Background()))
	}()

	return nil
}

func (h *Host) log(ctx context.Context, m api.Module, offset, byteCount uint32) {
	buf, ok := m.Memory().Read(offset, byteCount)
	if !ok {
		return
	}

	var taskLog *rpcv1.TaskLog
	if err := proto.Unmarshal(buf, taskLog); err != nil {
		h.Logger.Warn("failed to unmarshal task log", slog.String("err", err.Error()))
		return
	}

	rec := slog.NewRecord(taskLog.Time.AsTime(), slog.Level(taskLog.Level), taskLog.Message, -1)

	for i := range taskLog.Attributes {
		attr := taskLog.Attributes[i]
		rec.AddAttrs(slog.String(attr.Key, attr.Value))
	}

	_ = h.Logger.Handler().Handle(ctx, rec)
}