Moonshark/modules/registry.go

package modules

import (
	"embed"
	"fmt"
	"maps"
	"strings"

	"Moonshark/modules/crypto"
	"Moonshark/modules/fs"
	"Moonshark/modules/http"
	"Moonshark/modules/kv"
	"Moonshark/modules/sql"
	lua_string "Moonshark/modules/string+"

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

var Global *Registry

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

type Registry struct {
	modules       map[string]string
	globalModules map[string]string // globalName -> moduleSource
	goFuncs       map[string]luajit.GoFunction
}

func New() *Registry {
	r := &Registry{
		modules:       make(map[string]string),
		globalModules: make(map[string]string),
		goFuncs:       make(map[string]luajit.GoFunction),
	}

	maps.Copy(r.goFuncs, lua_string.GetFunctionList())
	maps.Copy(r.goFuncs, crypto.GetFunctionList())
	maps.Copy(r.goFuncs, fs.GetFunctionList())
	maps.Copy(r.goFuncs, http.GetFunctionList())
	maps.Copy(r.goFuncs, sql.GetFunctionList())
	maps.Copy(r.goFuncs, kv.GetFunctionList())

	r.loadEmbeddedModules()
	return r
}

func (r *Registry) loadEmbeddedModules() {
	dirs, _ := embeddedModules.ReadDir(".")

	for _, dir := range dirs {
		if !dir.IsDir() {
			continue
		}

		dirName := dir.Name()
		isGlobal := strings.HasSuffix(dirName, "+")

		var moduleName, globalName string
		if isGlobal {
			moduleName = strings.TrimSuffix(dirName, "+")
			globalName = moduleName
		} else {
			moduleName = dirName
		}

		modulePath := fmt.Sprintf("%s/%s.lua", dirName, moduleName)
		if source, err := embeddedModules.ReadFile(modulePath); err == nil {
			r.modules[moduleName] = string(source)
			if isGlobal {
				r.globalModules[globalName] = string(source)
			}
		}
	}
}

func (r *Registry) InstallInState(state *luajit.State) error {
	// Create moonshark global table with Go functions
	state.NewTable()
	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.SetGlobal("moonshark")

	// Auto-enhance all global modules
	for globalName, source := range r.globalModules {
		if err := r.enhanceGlobal(state, globalName, source); err != nil {
			return fmt.Errorf("failed to enhance %s global: %w", globalName, err)
		}
	}

	// Backup original require and install custom one
	state.GetGlobal("require")
	state.SetGlobal("_require_original")

	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")
		}

		// Return global if this module enhances a global
		if _, isGlobal := r.globalModules[moduleName]; isGlobal {
			s.GetGlobal(moduleName)
			return 1
		}

		// Check built-in modules
		if source, exists := r.modules[moduleName]; 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 {
				s.Pop(1)
				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
		}

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

func (r *Registry) enhanceGlobal(state *luajit.State, globalName, source string) error {
	// Execute the module - it directly modifies the global
	if err := state.LoadString(source); err != nil {
		return fmt.Errorf("failed to load %s module: %w", globalName, err)
	}
	if err := state.Call(0, 0); err != nil { // 0 results expected
		return fmt.Errorf("failed to execute %s module: %w", globalName, err)
	}

	return nil
}

func Initialize() error {
	Global = New()
	return nil
}