Moonshark/modules/registry.go

package modules

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

	"Moonshark/modules/crypto"
	"Moonshark/modules/fs"
	"Moonshark/modules/json"
	"Moonshark/modules/math"
	lua_string "Moonshark/modules/string"

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

// Global registry instance
var Global *Registry

//go:embed crypto/*.lua fs/*.lua json/*.lua math/*.lua string/*.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 from each module
	maps.Copy(r.goFuncs, json.GetFunctionList())
	maps.Copy(r.goFuncs, lua_string.GetFunctionList())
	maps.Copy(r.goFuncs, math.GetFunctionList())
	maps.Copy(r.goFuncs, fs.GetFunctionList())
	maps.Copy(r.goFuncs, crypto.GetFunctionList())

	return r
}

// LoadEmbeddedModules loads all .lua files from embedded filesystem
func (r *Registry) LoadEmbeddedModules() error {
	r.mutex.Lock()
	defer r.mutex.Unlock()

	// Load modules from subdirectories
	subdirs := []string{"crypto", "fs", "json", "math", "string"}
	for _, subdir := range subdirs {
		if err := r.loadModulesFromDir(subdir); err != nil {
			return err
		}
	}

	return nil
}

// loadModulesFromDir loads all .lua files from a specific directory
func (r *Registry) loadModulesFromDir(dir string) error {
	entries, err := embeddedModules.ReadDir(dir)
	if err != nil {
		return nil // Skip missing directories
	}

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

		fileName := strings.TrimSuffix(entry.Name(), ".lua")

		// If filename matches directory name (e.g., math/math.lua),
		// register as just the directory name for require("math")
		var moduleName string
		if fileName == dir {
			moduleName = dir
		} else {
			moduleName = dir + "/" + fileName
		}

		source, err := embeddedModules.ReadFile(filepath.Join(dir, 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()
}