package state import ( "fmt" "os" "path/filepath" "Moonshark/modules" luajit "git.sharkk.net/Sky/LuaJIT-to-Go" ) // State wraps luajit.State with enhanced functionality type State struct { *luajit.State initialized bool isWorker bool } // Config holds state initialization options type Config struct { OpenLibs bool InstallModules bool ScriptDir string IsWorker bool } // DefaultConfig returns default configuration func DefaultConfig() Config { return Config{ OpenLibs: true, InstallModules: true, IsWorker: false, } } // New creates a new enhanced Lua state func New(config ...Config) (*State, error) { cfg := DefaultConfig() if len(config) > 0 { cfg = config[0] } // Create base state baseState := luajit.New(cfg.OpenLibs) if baseState == nil { return nil, fmt.Errorf("failed to create Lua state") } state := &State{ State: baseState, isWorker: cfg.IsWorker, } // Set worker global flag if this is a worker state if cfg.IsWorker { state.PushBoolean(true) state.SetGlobal("__IS_WORKER") } // Install module system if requested if cfg.InstallModules { if err := state.initializeModules(); err != nil { state.Close() return nil, fmt.Errorf("failed to install modules: %w", err) } } // Set script directory if provided if cfg.ScriptDir != "" { if err := state.SetScriptDirectory(cfg.ScriptDir); err != nil { state.Close() return nil, fmt.Errorf("failed to set script directory: %w", err) } } state.initialized = true return state, nil } // NewFromScript creates a state configured for a specific script file func NewFromScript(scriptPath string, config ...Config) (*State, error) { cfg := DefaultConfig() if len(config) > 0 { cfg = config[0] } // Set script directory from file path scriptDir := filepath.Dir(scriptPath) absScriptDir, err := filepath.Abs(scriptDir) if err != nil { return nil, fmt.Errorf("failed to get absolute path: %w", err) } cfg.ScriptDir = absScriptDir return New(cfg) } // NewWorker creates a new worker state with __IS_WORKER flag set func NewWorker(config ...Config) (*State, error) { cfg := DefaultConfig() if len(config) > 0 { cfg = config[0] } cfg.IsWorker = true return New(cfg) } // NewWorkerFromScript creates a worker state configured for a specific script func NewWorkerFromScript(scriptPath string, config ...Config) (*State, error) { cfg := DefaultConfig() if len(config) > 0 { cfg = config[0] } cfg.IsWorker = true scriptDir := filepath.Dir(scriptPath) absScriptDir, err := filepath.Abs(scriptDir) if err != nil { return nil, fmt.Errorf("failed to get absolute path: %w", err) } cfg.ScriptDir = absScriptDir return New(cfg) } // initializeModules sets up the module system func (s *State) initializeModules() error { // Initialize global registry if needed if modules.Global == nil { if err := modules.Initialize(); err != nil { return err } } return modules.Global.InstallInState(s.State) } // SetScriptDirectory adds a directory to Lua's package.path func (s *State) SetScriptDirectory(dir string) error { packagePath := filepath.Join(dir, "?.lua") return s.AddPackagePath(packagePath) } // ExecuteFile compiles and runs a Lua script file func (s *State) ExecuteFile(scriptPath string) error { // Check if file exists if _, err := os.Stat(scriptPath); os.IsNotExist(err) { return fmt.Errorf("script file '%s' not found", scriptPath) } // Read script content scriptContent, err := os.ReadFile(scriptPath) if err != nil { return fmt.Errorf("failed to read script file '%s': %w", scriptPath, err) } return s.ExecuteString(string(scriptContent), scriptPath) } // ExecuteString compiles and runs Lua code func (s *State) ExecuteString(code, name string) error { // Use bytecode compilation for better performance bytecode, err := s.CompileBytecode(code, name) if err != nil { return fmt.Errorf("compilation error in '%s': %w", name, err) } if err := s.LoadAndRunBytecode(bytecode, name); err != nil { return fmt.Errorf("execution error in '%s': %w", name, err) } return nil } // ExecuteFileWithResults executes a script and returns results func (s *State) ExecuteFileWithResults(scriptPath string) ([]any, error) { if _, err := os.Stat(scriptPath); os.IsNotExist(err) { return nil, fmt.Errorf("script file '%s' not found", scriptPath) } scriptContent, err := os.ReadFile(scriptPath) if err != nil { return nil, fmt.Errorf("failed to read script file '%s': %w", scriptPath, err) } return s.ExecuteStringWithResults(string(scriptContent), scriptPath) } // ExecuteStringWithResults executes code and returns all results func (s *State) ExecuteStringWithResults(code, name string) ([]any, error) { baseTop := s.GetTop() defer s.SetTop(baseTop) nresults, err := s.Execute(code) if err != nil { return nil, fmt.Errorf("execution error in '%s': %w", name, err) } if nresults == 0 { return nil, nil } results := make([]any, nresults) for i := 0; i < nresults; i++ { val, err := s.ToValue(baseTop + i + 1) if err != nil { results[i] = nil } else { results[i] = val } } return results, nil } // IsInitialized returns whether the state has been properly initialized func (s *State) IsInitialized() bool { return s.initialized } // IsWorker returns whether this state is a worker state func (s *State) IsWorker() bool { return s.isWorker } // Close cleans up the state and releases resources func (s *State) Close() { if s.State != nil { s.Cleanup() s.State.Close() s.State = nil } s.initialized = false } // Quick creates a minimal state for one-off script execution func Quick() (*State, error) { return New(Config{ OpenLibs: true, InstallModules: true, }) } // QuickExecute creates a state, executes code, and cleans up func QuickExecute(code, name string) error { state, err := Quick() if err != nil { return err } defer state.Close() return state.ExecuteString(code, name) } // QuickExecuteWithResults executes code and returns results func QuickExecuteWithResults(code, name string) ([]any, error) { state, err := Quick() if err != nil { return nil, err } defer state.Close() return state.ExecuteStringWithResults(code, name) }