Moonshark/core/runner/ModuleLoader.go

package runner

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

	"Moonshark/core/utils/logger"

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

// ModuleConfig holds configuration for Lua's module loading system
type ModuleConfig struct {
	ScriptDir string   // Base directory for script being executed
	LibDirs   []string // Additional library directories
}

// ModuleLoader manages module loading and caching
type ModuleLoader struct {
	config        *ModuleConfig
	pathCache     map[string]string // Cache module paths for fast lookups
	bytecodeCache map[string][]byte // Cache of compiled bytecode
	debug         bool
	mu            sync.RWMutex
}

// NewModuleLoader creates a new module loader
func NewModuleLoader(config *ModuleConfig) *ModuleLoader {
	if config == nil {
		config = &ModuleConfig{
			ScriptDir: "",
			LibDirs:   []string{},
		}
	}

	return &ModuleLoader{
		config:        config,
		pathCache:     make(map[string]string),
		bytecodeCache: make(map[string][]byte),
		debug:         false,
	}
}

// EnableDebug turns on debug logging
func (l *ModuleLoader) EnableDebug() {
	l.debug = true
}

// SetScriptDir sets the script directory
func (l *ModuleLoader) SetScriptDir(dir string) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.config.ScriptDir = dir
}

// debugLog logs a message if debug mode is enabled
func (l *ModuleLoader) debugLog(format string, args ...interface{}) {
	if l.debug {
		logger.Debug("ModuleLoader "+format, args...)
	}
}

// SetupRequire configures the require system in a Lua state
func (l *ModuleLoader) SetupRequire(state *luajit.State) error {
	l.mu.RLock()
	defer l.mu.RUnlock()

	// Initialize our module registry in Lua
	err := state.DoString(`
		-- Initialize global module registry
		__module_paths = {}

		-- Setup fast module loading system
		__module_bytecode = {}

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

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

	if err != nil {
		return err
	}

	// Set up package.path based on search paths
	paths := l.getSearchPaths()
	pathStr := strings.Join(paths, ";")
	escapedPathStr := escapeLuaString(pathStr)

	return state.DoString(`package.path = "` + escapedPathStr + `"`)
}

// getSearchPaths returns a list of Lua search paths
func (l *ModuleLoader) getSearchPaths() []string {
	absPaths := []string{}
	seen := map[string]bool{}

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

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

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

	return absPaths
}

// PreloadModules preloads modules from library directories
func (l *ModuleLoader) PreloadModules(state *luajit.State) error {
	l.mu.Lock()
	defer l.mu.Unlock()

	// Reset caches
	l.pathCache = make(map[string]string)
	l.bytecodeCache = make(map[string][]byte)

	// Reset module registry in Lua
	if err := state.DoString(`
		-- Reset module registry
		__module_paths = {}
		__module_bytecode = {}
		__ready_modules = {}

		-- 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 = {}
	`); err != nil {
		return err
	}

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

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

		l.debugLog("Scanning directory: %s", absDir)

		// 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 from path
			relPath, err := filepath.Rel(absDir, path)
			if err != nil || strings.HasPrefix(relPath, "..") {
				return nil
			}

			// Convert path to module name
			modName := strings.TrimSuffix(relPath, ".lua")
			modName = strings.ReplaceAll(modName, string(filepath.Separator), ".")

			l.debugLog("Found module: %s at %s", modName, path)

			// Register in our caches
			l.pathCache[modName] = path

			// Load file content
			content, err := os.ReadFile(path)
			if err != nil {
				l.debugLog("Failed to read module file: %v", err)
				return nil
			}

			// Compile to bytecode
			bytecode, err := state.CompileBytecode(string(content), path)
			if err != nil {
				l.debugLog("Failed to compile module: %v", err)
				return nil
			}

			// Cache bytecode
			l.bytecodeCache[modName] = bytecode

			// Register in Lua
			escapedPath := escapeLuaString(path)
			escapedName := escapeLuaString(modName)

			if err := state.DoString(`__module_paths["` + escapedName + `"] = "` + escapedPath + `"`); err != nil {
				return nil
			}

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

			// Add to package.preload
			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(`
		-- Setup environment-aware require function
		function __setup_require(env)
			-- Create require function specific to this environment
			env.require = function(modname)
				-- Check if already loaded
				if package.loaded[modname] then
					return package.loaded[modname]
				end

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

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

						package.loaded[modname] = result
						return result
					end
				end

				-- 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

				-- Full path search as last resort
				local errors = {}
				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(errors, "\tno file '" .. file_path .. "'")
				end

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

			return env
		end
	`)
}

// GetModuleByPath finds the module name for a file path
func (l *ModuleLoader) GetModuleByPath(path string) (string, bool) {
	l.mu.RLock()
	defer l.mu.RUnlock()

	// Clean path for proper comparison
	path = filepath.Clean(path)

	// Try direct lookup from cache
	for modName, modPath := range l.pathCache {
		if modPath == path {
			return modName, true
		}
	}

	// Try to find by relative path from lib dirs
	for _, dir := range l.config.LibDirs {
		absDir, err := filepath.Abs(dir)
		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), ".")
			return modName, true
		}
	}

	return "", false
}

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