Moonshark/core/runner/ModuleLoader.go

package runner

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

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

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

// ModuleInfo stores information about a loaded module
type ModuleInfo struct {
	Name     string
	Path     string
	IsCore   bool
	Bytecode []byte
}

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

// ModuleRegistry keeps track of Lua modules for file watching
type ModuleRegistry struct {
	// Maps file paths to module names
	pathToModule sync.Map
	// Maps module names to file paths
	moduleToPath sync.Map
}

// NewModuleRegistry creates a new module registry
func NewModuleRegistry() *ModuleRegistry {
	return &ModuleRegistry{}
}

// Register adds a module path to the registry
func (r *ModuleRegistry) Register(path string, name string) {
	r.pathToModule.Store(path, name)
	r.moduleToPath.Store(name, path)
}

// GetModuleName retrieves a module name by path
func (r *ModuleRegistry) GetModuleName(path string) (string, bool) {
	value, ok := r.pathToModule.Load(path)
	if !ok {
		return "", false
	}
	return value.(string), true
}

// GetModulePath retrieves a path by module name
func (r *ModuleRegistry) GetModulePath(name string) (string, bool) {
	value, ok := r.moduleToPath.Load(name)
	if !ok {
		return "", false
	}
	return value.(string), true
}

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

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

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

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

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

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

		// 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), ".")

			// Register in our caches
			l.pathCache[modName] = path
			l.registry.Register(path, modName)

			// Load file content
			content, err := os.ReadFile(path)
			if err != nil {
				return nil
			}

			// Compile to bytecode
			bytecode, err := state.CompileBytecode(string(content), path)
			if err != nil {
				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 registry
	modName, found := l.registry.GetModuleName(path)
	if found {
		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
}

// ReloadModule reloads a module from disk
func (l *ModuleLoader) ReloadModule(state *luajit.State, name string) (bool, error) {
	l.mu.Lock()
	defer l.mu.Unlock()

	// Get module path
	path, ok := l.registry.GetModulePath(name)
	if !ok {
		for modName, modPath := range l.pathCache {
			if modName == name {
				path = modPath
				ok = true
				break
			}
		}
	}

	if !ok || path == "" {
		return false, nil
	}

	// Invalidate module in Lua
	err := state.DoString(`
		package.loaded["` + name + `"] = nil
		__ready_modules["` + name + `"] = nil
		if package.preload then
			package.preload["` + name + `"] = nil
		end
	`)

	if err != nil {
		return false, err
	}

	// Check if file still exists
	if _, err := os.Stat(path); os.IsNotExist(err) {
		// File was deleted, just invalidate
		delete(l.pathCache, name)
		delete(l.bytecodeCache, name)
		l.registry.moduleToPath.Delete(name)
		l.registry.pathToModule.Delete(path)
		return true, nil
	}

	// Read updated file
	content, err := os.ReadFile(path)
	if err != nil {
		return false, err
	}

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

	// Update cache
	l.bytecodeCache[name] = bytecode

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

	// Update preload
	luaCode := `
		local modname = "` + name + `"
		package.loaded[modname] = nil
		package.preload[modname] = ...
		__ready_modules[modname] = true
	`

	if err := state.DoString(luaCode); err != nil {
		state.Pop(1) // Remove chunk from stack
		return false, err
	}

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

// ResetModules clears non-core modules from package.loaded
func (l *ModuleLoader) 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
	`)
}

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