Moonshark/core/runner/Require.go

package runner

import (
	"os"
	"path/filepath"
	"strings"
	"sync"

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

// RequireConfig holds configuration for Lua's require mechanism
type RequireConfig struct {
	ScriptDir string   // Base directory for script being executed
	LibDirs   []string // Additional library directories
}

// NativeModuleLoader uses Lua's native package.loaded as the cache
type NativeModuleLoader struct {
	registry *ModuleRegistry
	config   *RequireConfig
	mu       sync.RWMutex
}

// NewNativeModuleLoader creates a new native module loader
func NewNativeModuleLoader(config *RequireConfig) *NativeModuleLoader {
	return &NativeModuleLoader{
		registry: NewModuleRegistry(),
		config:   config,
	}
}

// escapeLuaString escapes special characters in a string for Lua
func escapeLuaString(s string) string {
	replacer := strings.NewReplacer(
		"\\", "\\\\",
		"\"", "\\\"",
		"\n", "\\n",
		"\r", "\\r",
		"\t", "\\t",
	)
	return replacer.Replace(s)
}

// SetupRequire configures the require system
func (l *NativeModuleLoader) SetupRequire(state *luajit.State) error {
	// Initialize our module registry in Lua
	return state.DoString(`
		-- Initialize global module registry
		__module_paths = {}

		-- Setup fast module loading system
		__module_results = {}

		-- Create module preload table
		package.preload = package.preload or {}

		-- Setup module loader registry
		__ready_modules = {}
	`)
}

// PreloadAllModules fully preloads modules for maximum performance
func (l *NativeModuleLoader) PreloadAllModules(state *luajit.State) error {
	l.mu.Lock()
	defer l.mu.Unlock()

	// Reset registry
	l.registry = NewModuleRegistry()

	// Reset preloaded modules in Lua
	if err := state.DoString(`
		-- Reset module registry
		__module_paths = {}
		__module_results = {}

		-- Clear non-core modules from package.loaded
		local core_modules = {
			string = true, table = true, math = true, os = true,
			package = true, io = true, coroutine = true, debug = true, _G = true
		}

		for name in pairs(package.loaded) do
			if not core_modules[name] then
				package.loaded[name] = nil
			end
		end

		-- Reset preload table
		package.preload = package.preload or {}
		for name in pairs(package.preload) do
			package.preload[name] = nil
		end

		-- Reset ready modules
		__ready_modules = {}
	`); err != nil {
		return err
	}

	// Set up paths for require
	absPaths := []string{}
	pathsMap := map[string]bool{}

	// Add script directory (absolute path)
	if l.config.ScriptDir != "" {
		absPath, err := filepath.Abs(l.config.ScriptDir)
		if err == nil && !pathsMap[absPath] {
			absPaths = append(absPaths, filepath.Join(absPath, "?.lua"))
			pathsMap[absPath] = true
		}
	}

	// Add lib directories (absolute paths)
	for _, dir := range l.config.LibDirs {
		if dir == "" {
			continue
		}

		absPath, err := filepath.Abs(dir)
		if err == nil && !pathsMap[absPath] {
			absPaths = append(absPaths, filepath.Join(absPath, "?.lua"))
			pathsMap[absPath] = true
		}
	}

	// Set package.path
	escapedPathStr := escapeLuaString(strings.Join(absPaths, ";"))
	if err := state.DoString(`package.path = "` + escapedPathStr + `"`); err != nil {
		return err
	}

	// Process and preload all modules from lib directories
	for _, dir := range l.config.LibDirs {
		if dir == "" {
			continue
		}

		absDir, err := filepath.Abs(dir)
		if err != nil {
			continue
		}

		// Find all Lua files
		err = filepath.Walk(absDir, func(path string, info os.FileInfo, err error) error {
			if err != nil || info.IsDir() || !strings.HasSuffix(path, ".lua") {
				return nil
			}

			// Get module name
			relPath, err := filepath.Rel(absDir, path)
			if err != nil || strings.HasPrefix(relPath, "..") {
				return nil
			}

			modName := strings.TrimSuffix(relPath, ".lua")
			modName = strings.ReplaceAll(modName, string(filepath.Separator), ".")

			// Register module path
			l.registry.Register(path, modName)

			// Register path in Lua
			escapedPath := escapeLuaString(path)
			escapedName := escapeLuaString(modName)
			if err := state.DoString(`__module_paths["` + escapedName + `"] = "` + escapedPath + `"`); err != nil {
				return nil
			}

			// Compile the module
			content, err := os.ReadFile(path)
			if err != nil {
				return nil
			}

			// Precompile bytecode
			bytecode, err := state.CompileBytecode(string(content), path)
			if err != nil {
				return nil
			}

			// Load bytecode
			if err := state.LoadBytecode(bytecode, path); err != nil {
				return nil
			}

			// Store in package.preload for fast loading
			// We use string concat for efficiency (no string.format overhead)
			luaCode := `
				local modname = "` + escapedName + `"
				local chunk = ...
				package.preload[modname] = chunk
				__ready_modules[modname] = true
			`
			if err := state.DoString(luaCode); err != nil {
				state.Pop(1) // Remove chunk from stack
				return nil
			}

			state.Pop(1) // Remove chunk from stack
			return nil
		})

		if err != nil {
			return err
		}
	}

	// Install optimized require implementation
	return state.DoString(`
		-- Ultra-fast module loader
		function __fast_require(env, modname)
			-- 1. Check already loaded modules
			if package.loaded[modname] then
				return package.loaded[modname]
			end

			-- 2. Check preloaded chunks
			if __ready_modules[modname] then
				local loader = package.preload[modname]
				if loader then
					-- Set environment
					setfenv(loader, env)

					-- Execute and store result
					local result = loader()
					if result == nil then
						result = true
					end

					-- Cache in shared registry
					package.loaded[modname] = result
					return result
				end
			end

			-- 3. Direct file load as fallback
			if __module_paths[modname] then
				local path = __module_paths[modname]
				local chunk, err = loadfile(path)
				if chunk then
					setfenv(chunk, env)
					local result = chunk()
					if result == nil then
						result = true
					end
					package.loaded[modname] = result
					return result
				end
			end

			-- 4. Full path search as last resort
			local err_msgs = {}
			for path in package.path:gmatch("[^;]+") do
				local file_path = path:gsub("?", modname:gsub("%.", "/"))
				local chunk, err = loadfile(file_path)
				if chunk then
					setfenv(chunk, env)
					local result = chunk()
					if result == nil then
						result = true
					end
					package.loaded[modname] = result
					return result
				end
				table.insert(err_msgs, "no file '" .. file_path .. "'")
			end

			error("module '" .. modname .. "' not found:\n" .. table.concat(err_msgs, "\n"), 2)
		end

		-- Install require factory
		function __setup_require(env)
			-- Create highly optimized require with closure
			env.require = function(modname)
				return __fast_require(env, modname)
			end
			return env
		end
	`)
}

// NotifyFileChanged invalidates modules when files change
func (l *NativeModuleLoader) NotifyFileChanged(state *luajit.State, path string) bool {
	path = filepath.Clean(path)

	// Get module name from registry
	modName, found := l.registry.GetModuleName(path)
	if !found {
		// Try to find by path for lib dirs
		for _, libDir := range l.config.LibDirs {
			absDir, err := filepath.Abs(libDir)
			if err != nil {
				continue
			}

			relPath, err := filepath.Rel(absDir, path)
			if err != nil || strings.HasPrefix(relPath, "..") {
				continue
			}

			if strings.HasSuffix(relPath, ".lua") {
				modName = strings.TrimSuffix(relPath, ".lua")
				modName = strings.ReplaceAll(modName, string(filepath.Separator), ".")
				found = true
				break
			}
		}
	}

	if !found {
		return false
	}

	// Check if this is a core module
	coreName, isCoreModule := GlobalRegistry.MatchModuleName(modName)

	// Invalidate module in Lua
	escapedName := escapeLuaString(modName)
	invalidateCode := `
		package.loaded["` + escapedName + `"] = nil
		__ready_modules["` + escapedName + `"] = nil
		if package.preload then
			package.preload["` + escapedName + `"] = nil
		end
	`
	if err := state.DoString(invalidateCode); err != nil {
		return false
	}

	// For core modules, reinitialize them
	if isCoreModule {
		if err := GlobalRegistry.InitializeModule(state, coreName); err != nil {
			return false
		}
		return true
	}

	// For regular modules, update bytecode if the file still exists
	content, err := os.ReadFile(path)
	if err != nil {
		// File might have been deleted - just invalidate
		return true
	}

	// Recompile module
	bytecode, err := state.CompileBytecode(string(content), path)
	if err != nil {
		// Invalid Lua - just invalidate
		return true
	}

	// Load bytecode
	if err := state.LoadBytecode(bytecode, path); err != nil {
		// Unable to load - just invalidate
		return true
	}

	// Update preload with new chunk for regular modules
	luaCode := `
		-- Update module in package.preload and clear loaded
		package.loaded["` + escapedName + `"] = nil
		package.preload["` + escapedName + `"] = ...
		__ready_modules["` + escapedName + `"] = true
	`
	if err := state.DoString(luaCode); err != nil {
		state.Pop(1) // Remove chunk from stack
		return false
	}

	state.Pop(1) // Remove chunk from stack
	return true
}

// ResetModules clears all non-core modules
func (l *NativeModuleLoader) ResetModules(state *luajit.State) error {
	return state.DoString(`
		local core_modules = {
			string = true, table = true, math = true, os = true,
			package = true, io = true, coroutine = true, debug = true, _G = true
		}

		for name in pairs(package.loaded) do
			if not core_modules[name] then
				package.loaded[name] = nil
			end
		end
	`)
}