Moonshark/modules/registry.go

package registry

import (
	"embed"
	"fmt"
	"maps"
	"path/filepath"
	"strings"
	"sync"

	"Moonshark/functions"

	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
)

// Global registry instance
var Global *Registry

//go:embed modules/*.lua
var embeddedModules embed.FS

// Registry manages all Lua modules and Go functions
type Registry struct {
	modules map[string]string
	goFuncs map[string]luajit.GoFunction
	mutex   sync.RWMutex
}

// New creates a new registry with all Go functions loaded
func New() *Registry {
	r := &Registry{
		modules: make(map[string]string),
		goFuncs: make(map[string]luajit.GoFunction),
	}

	// Load all Go functions
	maps.Copy(r.goFuncs, functions.GetJSONFunctions())
	maps.Copy(r.goFuncs, functions.GetStringFunctions())
	maps.Copy(r.goFuncs, functions.GetMathFunctions())
	maps.Copy(r.goFuncs, functions.GetFSFunctions())
	maps.Copy(r.goFuncs, functions.GetCryptoFunctions())

	return r
}

// LoadEmbeddedModules loads all .lua files from embedded filesystem
func (r *Registry) LoadEmbeddedModules() error {
	entries, err := embeddedModules.ReadDir("modules")
	if err != nil {
		return fmt.Errorf("failed to read modules directory: %w", err)
	}

	r.mutex.Lock()
	defer r.mutex.Unlock()

	for _, entry := range entries {
		if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".lua") {
			continue
		}

		moduleName := strings.TrimSuffix(entry.Name(), ".lua")
		source, err := embeddedModules.ReadFile(filepath.Join("modules", entry.Name()))
		if err != nil {
			return fmt.Errorf("failed to read module %s: %w", moduleName, err)
		}

		r.modules[moduleName] = string(source)
	}

	return nil
}

// InstallInState sets up the complete module system in a Lua state
func (r *Registry) InstallInState(state *luajit.State) error {
	r.mutex.RLock()
	defer r.mutex.RUnlock()

	// Create moonshark global table
	state.NewTable()
	state.SetGlobal("moonshark")

	// Install Go functions
	state.GetGlobal("moonshark")
	for name, fn := range r.goFuncs {
		if err := state.PushGoFunction(fn); err != nil {
			return fmt.Errorf("failed to register Go function '%s': %w", name, err)
		}
		state.SetField(-2, name)
	}
	state.Pop(1) // Remove moonshark table

	// Backup original require
	state.GetGlobal("require")
	state.SetGlobal("_require_original")

	// Install custom require function
	return state.RegisterGoFunction("require", func(s *luajit.State) int {
		if err := s.CheckMinArgs(1); err != nil {
			return s.PushError("require: %v", err)
		}

		moduleName, err := s.SafeToString(1)
		if err != nil {
			return s.PushError("require: module name must be a string")
		}

		// Check built-in modules first
		r.mutex.RLock()
		source, exists := r.modules[moduleName]
		r.mutex.RUnlock()

		if exists {
			if err := s.LoadString(source); err != nil {
				return s.PushError("require: failed to load module '%s': %v", moduleName, err)
			}
			if err := s.Call(0, 1); err != nil {
				return s.PushError("require: failed to execute module '%s': %v", moduleName, err)
			}
			return 1
		}

		// Fall back to original require
		s.GetGlobal("_require_original")
		if s.IsFunction(-1) {
			s.PushString(moduleName)
			if err := s.Call(1, 1); err != nil {
				return s.PushError("require: %v", err)
			}
			return 1
		}

		return s.PushError("require: module '%s' not found", moduleName)
	})
}

// Initialize sets up the global registry with all modules loaded
func Initialize() error {
	Global = New()
	return Global.LoadEmbeddedModules()
}

// GetGoFunctions returns all Go functions
func (r *Registry) GetGoFunctions() map[string]luajit.GoFunction {
	r.mutex.RLock()
	defer r.mutex.RUnlock()

	result := make(map[string]luajit.GoFunction, len(r.goFuncs))
	for k, v := range r.goFuncs {
		result[k] = v
	}
	return result
}