package state

import (
	"context"

	"flux/core"
)

// TODO: currently using Golang Map for go through the DAG pipline
// This should be use index map index map[uint64]uint32  Hash(KeyGroup + StateName + Key) -> Offset
// with 1GB of []byte allocation
// This should be optimize with ZERO GC Scanning, Memory Alignment

type MemoryStateBackend struct {
	maxParallelism   int
	currentKey       string
	currentGroupKey  int
	currentNamespace any
	// 4D structure StateName -> KeyGroup -> Key -> Namespace -> Value
	states map[string]map[int]map[string]map[any]any
}

func NewMemoryStateBackend(maxParallelism int) *MemoryStateBackend {
	return &MemoryStateBackend{
		maxParallelism: maxParallelism,
		states:         make(map[string]map[int]map[string]map[any]any),
		// default is void namespace
		currentNamespace: core.VoidNamespace{},
	}
}

func (ms *MemoryStateBackend) GetCurrentKey() string {
	return ms.currentKey
}

func (ms *MemoryStateBackend) SetCurrentKey(key string) error {
	ms.currentKey = key
	ms.currentGroupKey = core.AssignToKeyGroup(key, ms.maxParallelism)
	return nil
}

func (ms *MemoryStateBackend) GetValueState(ctx context.Context, statename string) (core.ValueState[any], error) {
	// Assign if state not existed
	
	if _, ok := ms.states[statename]; !ok {
		// create state name map if not existed 
		ms.states[statename] = make(map[int]map[string]map[any]any)
	}

	return &memoryStateValue{backend: ms, stateName: statename}, nil
}

func (ms *MemoryStateBackend) GetCurrentNamespace() any {
	return ms.currentNamespace
}

func (ms *MemoryStateBackend) SetCurrentNamespace(namespace any) error {
	ms.currentNamespace = namespace
	return nil
}