Moonshark/modules/kv/kv.go

package kv

import (
	"bufio"
	"fmt"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
	"github.com/goccy/go-json"
)

var (
	stores = make(map[string]*Store)
	mutex  sync.RWMutex
)

type Store struct {
	data     map[string]string
	expires  map[string]int64
	filename string
	mutex    sync.RWMutex
}

func GetFunctionList() map[string]luajit.GoFunction {
	return map[string]luajit.GoFunction{
		"kv_open":            kv_open,
		"kv_get":             kv_get,
		"kv_set":             kv_set,
		"kv_delete":          kv_delete,
		"kv_clear":           kv_clear,
		"kv_has":             kv_has,
		"kv_size":            kv_size,
		"kv_keys":            kv_keys,
		"kv_values":          kv_values,
		"kv_save":            kv_save,
		"kv_close":           kv_close,
		"kv_increment":       kv_increment,
		"kv_append":          kv_append,
		"kv_expire":          kv_expire,
		"kv_cleanup_expired": kv_cleanup_expired,
	}
}

// kv_open(name, filename) -> boolean
func kv_open(s *luajit.State) int {
	name := s.ToString(1)
	filename := s.ToString(2)

	if name == "" {
		s.PushBoolean(false)
		return 1
	}

	mutex.Lock()
	defer mutex.Unlock()

	if store, exists := stores[name]; exists {
		if filename != "" && store.filename != filename {
			store.filename = filename
		}
		s.PushBoolean(true)
		return 1
	}

	store := &Store{
		data:     make(map[string]string),
		expires:  make(map[string]int64),
		filename: filename,
	}

	if filename != "" {
		store.load()
	}

	stores[name] = store
	s.PushBoolean(true)
	return 1
}

// kv_get(name, key, default) -> value or default
func kv_get(s *luajit.State) int {
	name := s.ToString(1)
	key := s.ToString(2)
	hasDefault := s.GetTop() >= 3

	mutex.RLock()
	store, exists := stores[name]
	mutex.RUnlock()

	if !exists {
		if hasDefault {
			s.PushCopy(3)
		} else {
			s.PushNil()
		}
		return 1
	}

	store.mutex.RLock()
	value, found := store.data[key]
	store.mutex.RUnlock()

	if found {
		s.PushString(value)
	} else if hasDefault {
		s.PushCopy(3)
	} else {
		s.PushNil()
	}
	return 1
}

// kv_set(name, key, value) -> boolean
func kv_set(s *luajit.State) int {
	name := s.ToString(1)
	key := s.ToString(2)
	value := s.ToString(3)

	mutex.RLock()
	store, exists := stores[name]
	mutex.RUnlock()

	if !exists {
		s.PushBoolean(false)
		return 1
	}

	store.mutex.Lock()
	store.data[key] = value
	store.mutex.Unlock()

	s.PushBoolean(true)
	return 1
}

// kv_delete(name, key) -> boolean
func kv_delete(s *luajit.State) int {
	name := s.ToString(1)
	key := s.ToString(2)

	mutex.RLock()
	store, exists := stores[name]
	mutex.RUnlock()

	if !exists {
		s.PushBoolean(false)
		return 1
	}

	store.mutex.Lock()
	_, existed := store.data[key]
	delete(store.data, key)
	delete(store.expires, key)
	store.mutex.Unlock()

	s.PushBoolean(existed)
	return 1
}

// kv_clear(name) -> boolean
func kv_clear(s *luajit.State) int {
	name := s.ToString(1)

	mutex.RLock()
	store, exists := stores[name]
	mutex.RUnlock()

	if !exists {
		s.PushBoolean(false)
		return 1
	}

	store.mutex.Lock()
	store.data = make(map[string]string)
	store.expires = make(map[string]int64)
	store.mutex.Unlock()

	s.PushBoolean(true)
	return 1
}

// kv_has(name, key) -> boolean
func kv_has(s *luajit.State) int {
	name := s.ToString(1)
	key := s.ToString(2)

	mutex.RLock()
	store, exists := stores[name]
	mutex.RUnlock()

	if !exists {
		s.PushBoolean(false)
		return 1
	}

	store.mutex.RLock()
	_, found := store.data[key]
	store.mutex.RUnlock()

	s.PushBoolean(found)
	return 1
}

// kv_size(name) -> number
func kv_size(s *luajit.State) int {
	name := s.ToString(1)

	mutex.RLock()
	store, exists := stores[name]
	mutex.RUnlock()

	if !exists {
		s.PushNumber(0)
		return 1
	}

	store.mutex.RLock()
	size := len(store.data)
	store.mutex.RUnlock()

	s.PushNumber(float64(size))
	return 1
}

// kv_keys(name) -> table
func kv_keys(s *luajit.State) int {
	name := s.ToString(1)

	mutex.RLock()
	store, exists := stores[name]
	mutex.RUnlock()

	if !exists {
		s.NewTable()
		return 1
	}

	store.mutex.RLock()
	keys := make([]string, 0, len(store.data))
	for k := range store.data {
		keys = append(keys, k)
	}
	store.mutex.RUnlock()

	s.PushValue(keys)
	return 1
}

// kv_values(name) -> table
func kv_values(s *luajit.State) int {
	name := s.ToString(1)

	mutex.RLock()
	store, exists := stores[name]
	mutex.RUnlock()

	if !exists {
		s.NewTable()
		return 1
	}

	store.mutex.RLock()
	values := make([]string, 0, len(store.data))
	for _, v := range store.data {
		values = append(values, v)
	}
	store.mutex.RUnlock()

	s.PushValue(values)
	return 1
}

// kv_save(name) -> boolean
func kv_save(s *luajit.State) int {
	name := s.ToString(1)

	mutex.RLock()
	store, exists := stores[name]
	mutex.RUnlock()

	if !exists || store.filename == "" {
		s.PushBoolean(false)
		return 1
	}

	err := store.save()
	s.PushBoolean(err == nil)
	return 1
}

// kv_close(name) -> boolean
func kv_close(s *luajit.State) int {
	name := s.ToString(1)

	mutex.Lock()
	defer mutex.Unlock()

	store, exists := stores[name]
	if !exists {
		s.PushBoolean(false)
		return 1
	}

	if store.filename != "" {
		store.save()
	}

	delete(stores, name)
	s.PushBoolean(true)
	return 1
}

// kv_increment(name, key, delta) -> number
func kv_increment(s *luajit.State) int {
	name := s.ToString(1)
	key := s.ToString(2)
	delta := 1.0
	if s.GetTop() >= 3 {
		delta = s.ToNumber(3)
	}

	mutex.RLock()
	store, exists := stores[name]
	mutex.RUnlock()

	if !exists {
		s.PushNumber(0)
		return 1
	}

	store.mutex.Lock()
	current, _ := strconv.ParseFloat(store.data[key], 64)
	newValue := current + delta
	store.data[key] = strconv.FormatFloat(newValue, 'g', -1, 64)
	store.mutex.Unlock()

	s.PushNumber(newValue)
	return 1
}

// kv_append(name, key, value, separator) -> boolean
func kv_append(s *luajit.State) int {
	name := s.ToString(1)
	key := s.ToString(2)
	value := s.ToString(3)
	separator := ""
	if s.GetTop() >= 4 {
		separator = s.ToString(4)
	}

	mutex.RLock()
	store, exists := stores[name]
	mutex.RUnlock()

	if !exists {
		s.PushBoolean(false)
		return 1
	}

	store.mutex.Lock()
	current := store.data[key]
	if current == "" {
		store.data[key] = value
	} else {
		store.data[key] = current + separator + value
	}
	store.mutex.Unlock()

	s.PushBoolean(true)
	return 1
}

// kv_expire(name, key, ttl) -> boolean
func kv_expire(s *luajit.State) int {
	name := s.ToString(1)
	key := s.ToString(2)
	ttl := s.ToNumber(3)

	mutex.RLock()
	store, exists := stores[name]
	mutex.RUnlock()

	if !exists {
		s.PushBoolean(false)
		return 1
	}

	store.mutex.Lock()
	store.expires[key] = time.Now().Unix() + int64(ttl)
	store.mutex.Unlock()

	s.PushBoolean(true)
	return 1
}

// kv_cleanup_expired(name) -> number
func kv_cleanup_expired(s *luajit.State) int {
	name := s.ToString(1)

	mutex.RLock()
	store, exists := stores[name]
	mutex.RUnlock()

	if !exists {
		s.PushNumber(0)
		return 1
	}

	currentTime := time.Now().Unix()
	deleted := 0

	store.mutex.Lock()
	for key, expireTime := range store.expires {
		if currentTime >= expireTime {
			delete(store.data, key)
			delete(store.expires, key)
			deleted++
		}
	}
	store.mutex.Unlock()

	s.PushNumber(float64(deleted))
	return 1
}

func (store *Store) load() error {
	if store.filename == "" {
		return nil
	}

	file, err := os.Open(store.filename)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}
	defer file.Close()

	if strings.HasSuffix(store.filename, ".json") {
		decoder := json.NewDecoder(file)
		return decoder.Decode(&store.data)
	}

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}

		parts := strings.SplitN(line, "=", 2)
		if len(parts) == 2 {
			store.data[parts[0]] = parts[1]
		}
	}

	return scanner.Err()
}

func (store *Store) save() error {
	if store.filename == "" {
		return fmt.Errorf("no filename specified")
	}

	file, err := os.Create(store.filename)
	if err != nil {
		return err
	}
	defer file.Close()

	store.mutex.RLock()
	defer store.mutex.RUnlock()

	if strings.HasSuffix(store.filename, ".json") {
		encoder := json.NewEncoder(file)
		encoder.SetIndent("", "\t")
		return encoder.Encode(store.data)
	}

	for key, value := range store.data {
		if _, err := fmt.Fprintf(file, "%s=%s\n", key, value); err != nil {
			return err
		}
	}

	return nil
}

// CloseAllStores saves and closes all open stores
func CloseAllStores() {
	mutex.Lock()
	defer mutex.Unlock()

	for name, store := range stores {
		if store.filename != "" {
			store.save()
		}
		delete(stores, name)
	}
}