Moonshark/state/state.go

package state

import (
	"fmt"
	"os"
	"path/filepath"

	"Moonshark/modules"
	"Moonshark/modules/http"
	"Moonshark/modules/kv"

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

// State wraps luajit.State with enhanced functionality
type State struct {
	*luajit.State
	initialized bool
	isWorker    bool
	scriptDir   string
}

// 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,
		scriptDir: cfg.ScriptDir,
	}

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

// Store main state initialization for worker replication
var (
	mainStateScriptDir  string
	mainStateScript     string
	mainStateScriptName string
)

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

	// Install modules first
	if err := modules.Global.InstallInState(s.State); err != nil {
		return err
	}

	return nil
}

// SetupStateCreator sets up the state creator after script is loaded
func (s *State) SetupStateCreator() {
	if s.isWorker {
		return
	}

	http.SetStateCreator(func() (*luajit.State, error) {
		cfg := DefaultConfig()
		cfg.IsWorker = true
		cfg.ScriptDir = mainStateScriptDir

		workerState, err := New(cfg)
		if err != nil {
			return nil, err
		}

		// Execute the same script as main state to get identical environment
		if mainStateScript != "" {
			if err := workerState.ExecuteString(mainStateScript, mainStateScriptName); err != nil {
				workerState.Close()
				return nil, fmt.Errorf("failed to execute script in worker: %w", err)
			}
		}

		return workerState.State, nil
	})
}

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

	// Store for worker replication if this is main state
	if !s.isWorker {
		mainStateScript = string(scriptContent)
		mainStateScriptName = scriptPath
		mainStateScriptDir = s.scriptDir
		// Set up state creator now that we have the script
		s.SetupStateCreator()
	}

	return s.ExecuteString(string(scriptContent), scriptPath)
}

// ExecuteString compiles and runs Lua code with bytecode caching
func (s *State) ExecuteString(code, name string) error {
	entry, err := CompileAndCache(s.State, code, name)
	if err != nil {
		return fmt.Errorf("compilation error in '%s': %w", name, err)
	}

	if err := s.LoadAndRunBytecode(entry.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 with bytecode caching
func (s *State) ExecuteStringWithResults(code, name string) ([]any, error) {
	baseTop := s.GetTop()
	defer s.SetTop(baseTop)

	entry, err := CompileAndCache(s.State, code, name)
	if err != nil {
		return nil, fmt.Errorf("compilation error in '%s': %w", name, err)
	}

	if err := s.LoadAndRunBytecodeWithResults(entry.Bytecode, name, -1); err != nil {
		return nil, fmt.Errorf("execution error in '%s': %w", name, err)
	}

	nresults := s.GetTop() - baseTop
	if nresults == 0 {
		return nil, nil
	}

	results := make([]any, nresults)
	for i := range nresults {
		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 {
		if !s.isWorker {
			kv.CloseAllStores()
		}

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