gomax/state/state.go

package state

import (
	"encoding/json"
	"fmt"

	"github.com/rs/zerolog/log"
)

var (
	ErrStateNotFound = fmt.Errorf("state not found")
	ErrKeyNotFound   = fmt.Errorf("state: key not found")
)

const ClearState State = "clear"

type State string

type Codec interface {
	Marshal(v any) ([]byte, error)
	Unmarshal(data []byte, v any) error
}

type JSONCodec struct{}

func (JSONCodec) Marshal(v any) ([]byte, error) {
	return json.Marshal(v)
}

func (JSONCodec) Unmarshal(d []byte, v any) error {
	return json.Unmarshal(d, v)
}

type Snapshot struct {
	ID    int               `json:"id"`
	State State             `json:"state"`
	Data  map[string][]byte `json:"data"`
}

type Context struct {
	id    int
	state State
	data  map[string][]byte
	codec Codec
}

type Option func(*Context)

func WithCodec(c Codec) Option { return func(sc *Context) { sc.codec = c } }

func New(id int, opts ...Option) *Context {
	sc := &Context{
		id:    id,
		data:  make(map[string][]byte),
		codec: JSONCodec{},
	}
	for _, o := range opts {
		o(sc)
	}
	return sc
}

func Restore(s Snapshot, opts ...Option) *Context {
	sc := New(s.ID, opts...)
	sc.state = s.State
	for k, v := range s.Data {
		sc.data[k] = v
	}
	return sc
}

func (sc *Context) Snapshot() Snapshot {
	data := make(map[string][]byte, len(sc.data))
	for k, v := range sc.data {
		data[k] = v
	}
	return Snapshot{ID: sc.id, State: sc.state, Data: data}
}

func (sc *Context) Set(state State) {
	log.Debug().
		Int("user_id", sc.id).
		Str("from", string(sc.state)).
		Str("to", string(state)).
		Msg("fsm state transition")
	sc.state = state
}

func (sc *Context) State() State { return sc.state }

func (sc *Context) Clear() {
	sc.state = ClearState
	sc.data = make(map[string][]byte)
}

func (sc *Context) Store(key string, val any) error {
	b, err := sc.codec.Marshal(val)
	if err != nil {
		return fmt.Errorf("sc.Store: failed to encode %q: %w", key, err)
	}
	sc.data[key] = b
	return nil
}

func (sc *Context) Delete(key string) {
	delete(sc.data, key)
}

func (sc *Context) Load(key string) Value {
	b, ok := sc.data[key]
	return Value{key: key, raw: b, ok: ok, codec: sc.codec}
}

type Value struct {
	key   string
	raw   []byte
	ok    bool
	codec Codec
}

func (v Value) Exists() bool { return v.ok }

func (v Value) Raw() []byte { return v.raw }

func (v Value) Decode(dst any) error {
	if !v.ok {
		return fmt.Errorf("%w: %q", ErrKeyNotFound, v.key)
	}
	if err := v.codec.Unmarshal(v.raw, dst); err != nil {
		return fmt.Errorf("state: decode %q as %T: %w", v.key, dst, err)
	}
	return nil
}

func (v Value) Str() (string, error) {
	var s string
	if err := v.Decode(&s); err != nil {
		return "", err
	}
	return s, nil
}
func (v Value) Int() (int, error) {
	var i int
	if err := v.Decode(&i); err != nil {
		return 0, err
	}
	return i, nil
}

func (v Value) Bool() (bool, error) {
	var b bool
	if err := v.Decode(&b); err != nil {
		return false, err
	}
	return b, nil
}
func (v Value) Float() (float64, error) {
	var f float64
	if err := v.Decode(&f); err != nil {
		return 0, err
	}
	return f, nil
}
func (v Value) Bytes() ([]byte, error) {
	var b []byte
	if err := v.Decode(&b); err != nil {
		return nil, err
	}
	return b, nil
}

func (v Value) StrOr(def string) string {
	if s, err := v.Str(); err == nil {
		return s
	}
	return def
}

func (v Value) IntOr(def int) int {
	if i, err := v.Int(); err == nil {
		return i
	}
	return def
}

func (v Value) BoolOr(def bool) bool {
	if b, err := v.Bool(); err == nil {
		return b
	}
	return def
}

func (v Value) FloatOr(def float64) float64 {
	if f, err := v.Float(); err == nil {
		return f
	}
	return def
}