package state

import (
	"encoding/json"
	"fmt"
	"github.com/abelith/gomax/v1/isolation"
	"strings"
)

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

type State struct {
	group string
	name  string
}

type Group struct {
	name   string
	states map[string]*State
}

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 {
	Scope isolation.Scope   `json:"id"`
	State string            `json:"state"`
	Data  map[string][]byte `json:"data"`
}

type Context struct {
	scope isolation.Scope
	state State
	data  map[string][]byte
	codec Codec
	clear bool
}

type Option func(*Context)

func NewStateGroup(name string) *Group {
	return &Group{
		name:   name,
		states: make(map[string]*State),
	}
}

func (g *Group) State(name string) State {
	state := g.states[name]

	if state != nil {
		return *state
	}

	if state = g.states[name]; state != nil {
		return *state
	}

	state = &State{
		group: g.name,
		name:  name,
	}

	g.states[name] = state
	return *state
}

func (s *State) String() string {
	return fmt.Sprintf("%s:%s", s.group, s.name)
}

func fromString(s string) State {
	pieces := strings.Split(s, ":")
	if len(pieces) < 2 {
		return State{}
	}

	return State{
		group: pieces[0],
		name:  pieces[1],
	}
}

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

func New(scope isolation.Scope, opts ...Option) *Context {
	sc := &Context{
		scope: scope,
		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.Scope, opts...)
	sc.state = fromString(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{Scope: sc.scope, State: sc.state.String(), Data: data}
}

func (sc *Context) Set(state State) {
	sc.state = state
}

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

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

func (sc *Context) IsEmpty() bool {
	return sc.state == State{} && sc.clear
}

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
}