Moonshark/core/runner/Runner.go

package runner

import (
	"Moonshark/core/utils/logger"
	"context"
	"errors"
	"path/filepath"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

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

// Common errors
var (
	ErrRunnerClosed  = errors.New("lua runner is closed")
	ErrInitFailed    = errors.New("initialization failed")
	ErrStateNotReady = errors.New("lua state not ready")
	ErrTimeout       = errors.New("operation timed out")
)

// RunnerOption defines a functional option for configuring the Runner
type RunnerOption func(*Runner)

// State wraps a Lua state with its sandbox
type State struct {
	L       *luajit.State // The Lua state
	sandbox *Sandbox      // Associated sandbox
	index   int           // Index for debugging
	inUse   bool          // Whether the state is currently in use
}

// Runner runs Lua scripts using a pool of Lua states
type Runner struct {
	states       []*State      // All states managed by this runner
	statePool    chan int      // Pool of available state indexes
	poolSize     int           // Size of the state pool
	moduleLoader *ModuleLoader // Module loader
	isRunning    atomic.Bool   // Whether the runner is active
	mu           sync.RWMutex  // Mutex for thread safety
	scriptDir    string        // Current script directory
}

// WithPoolSize sets the state pool size
func WithPoolSize(size int) RunnerOption {
	return func(r *Runner) {
		if size > 0 {
			r.poolSize = size
		}
	}
}

// WithLibDirs sets additional library directories
func WithLibDirs(dirs ...string) RunnerOption {
	return func(r *Runner) {
		if r.moduleLoader == nil {
			r.moduleLoader = NewModuleLoader(&ModuleConfig{
				LibDirs: dirs,
			})
		} else {
			r.moduleLoader.config.LibDirs = dirs
		}
	}
}

// NewRunner creates a new Runner with a pool of states
func NewRunner(options ...RunnerOption) (*Runner, error) {
	// Default configuration
	runner := &Runner{
		poolSize: runtime.GOMAXPROCS(0),
	}

	// Apply options
	for _, opt := range options {
		opt(runner)
	}

	// Set up module loader if not already initialized
	if runner.moduleLoader == nil {
		config := &ModuleConfig{
			ScriptDir: "",
			LibDirs:   []string{},
		}
		runner.moduleLoader = NewModuleLoader(config)
	}

	// Initialize states and pool
	runner.states = make([]*State, runner.poolSize)
	runner.statePool = make(chan int, runner.poolSize)

	// Create and initialize all states
	if err := runner.initializeStates(); err != nil {
		runner.Close() // Clean up already created states
		return nil, err
	}

	runner.isRunning.Store(true)
	return runner, nil
}

// initializeStates creates and initializes all states in the pool
func (r *Runner) initializeStates() error {
	logger.Server("Initializing %d states...", r.poolSize)

	for i := range r.poolSize {
		state, err := r.createState(i)
		if err != nil {
			return err
		}

		r.states[i] = state
		r.statePool <- i // Add index to the pool
	}

	return nil
}

// createState initializes a new Lua state
func (r *Runner) createState(index int) (*State, error) {
	verbose := index == 0
	if verbose {
		logger.ServerCont("Creating Lua state %d", index)
	}

	L := luajit.New()
	if L == nil {
		return nil, errors.New("failed to create Lua state")
	}

	sb := NewSandbox()

	// Set up sandbox
	if err := sb.Setup(L, verbose); err != nil {
		L.Cleanup()
		L.Close()
		return nil, ErrInitFailed
	}

	// Set up module loader
	if err := r.moduleLoader.SetupRequire(L); err != nil {
		L.Cleanup()
		L.Close()
		return nil, ErrInitFailed
	}

	// Preload modules
	if err := r.moduleLoader.PreloadModules(L); err != nil {
		L.Cleanup()
		L.Close()
		return nil, errors.New("failed to preload modules")
	}

	if verbose {
		logger.ServerCont("Lua state %d initialized successfully", index)
	}

	return &State{
		L:       L,
		sandbox: sb,
		index:   index,
		inUse:   false,
	}, nil
}

// Execute runs a script in a sandbox with context
func (r *Runner) Execute(ctx context.Context, bytecode []byte, execCtx *Context, scriptPath string) (*Response, error) {
	if !r.isRunning.Load() {
		return nil, ErrRunnerClosed
	}

	// Set script directory if provided
	if scriptPath != "" {
		r.mu.Lock()
		r.scriptDir = filepath.Dir(scriptPath)
		r.moduleLoader.SetScriptDir(r.scriptDir)
		r.mu.Unlock()
	}

	// Get a state from the pool
	var stateIndex int
	select {
	case stateIndex = <-r.statePool:
		// Got a state
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-time.After(5 * time.Second):
		return nil, ErrTimeout
	}

	// Get the actual state
	state := r.states[stateIndex]
	if state == nil {
		r.statePool <- stateIndex
		return nil, ErrStateNotReady
	}

	// Mark state as in use
	state.inUse = true

	// Ensure state is returned to pool when done
	defer func() {
		state.inUse = false
		if r.isRunning.Load() {
			select {
			case r.statePool <- stateIndex:
				// State returned to pool
			default:
				// Pool is full or closed
			}
		}
	}()

	// Execute in sandbox
	response, err := state.sandbox.Execute(state.L, bytecode, execCtx)
	if err != nil {
		return nil, err
	}

	return response, nil
}

// Run executes a Lua script with immediate context
func (r *Runner) Run(bytecode []byte, execCtx *Context, scriptPath string) (*Response, error) {
	return r.Execute(context.Background(), bytecode, execCtx, scriptPath)
}

// Close gracefully shuts down the Runner
func (r *Runner) Close() error {
	r.mu.Lock()
	defer r.mu.Unlock()

	if !r.isRunning.Load() {
		return ErrRunnerClosed
	}

	r.isRunning.Store(false)

	// Drain the state pool
	for {
		select {
		case <-r.statePool:
			// Drain one state
		default:
			// Pool is empty
			goto cleanup
		}
	}

cleanup:
	// Clean up all states
	for i, state := range r.states {
		if state != nil {
			state.L.Cleanup()
			state.L.Close()
			r.states[i] = nil
		}
	}

	logger.Debug("Runner closed")
	return nil
}

// RefreshStates rebuilds all states in the pool
func (r *Runner) RefreshStates() error {
	r.mu.Lock()
	defer r.mu.Unlock()

	if !r.isRunning.Load() {
		return ErrRunnerClosed
	}

	logger.Server("Runner is refreshing all states...")

	// Drain all states from the pool
	for {
		select {
		case <-r.statePool:
			// Drain one state
		default:
			// Pool is empty
			goto cleanup
		}
	}

cleanup:
	// Destroy all existing states
	for i, state := range r.states {
		if state != nil {
			if state.inUse {
				logger.WarningCont("Attempting to refresh state %d that is in use", i)
			}
			state.L.Cleanup()
			state.L.Close()
			r.states[i] = nil
		}
	}

	// Reinitialize all states
	if err := r.initializeStates(); err != nil {
		return err
	}

	logger.ServerCont("All states refreshed successfully")
	return nil
}

// NotifyFileChanged alerts the runner about file changes
func (r *Runner) NotifyFileChanged(filePath string) bool {
	logger.Debug("Runner has been notified of a file change...")
	logger.Debug("%s", filePath)

	// Check if it's a module file
	module, isModule := r.moduleLoader.GetModuleByPath(filePath)
	if isModule {
		logger.DebugCont("File is a module: %s", module)
		return r.RefreshModule(module)
	}

	// For non-module files, refresh all states
	if err := r.RefreshStates(); err != nil {
		logger.DebugCont("Failed to refresh states: %v", err)
		return false
	}

	return true
}

// RefreshModule refreshes a specific module across all states
func (r *Runner) RefreshModule(moduleName string) bool {
	r.mu.RLock()
	defer r.mu.RUnlock()

	if !r.isRunning.Load() {
		return false
	}

	logger.DebugCont("Refreshing module: %s", moduleName)

	success := true
	for _, state := range r.states {
		if state == nil || state.inUse {
			continue
		}

		// Invalidate module in Lua
		if err := state.L.DoString(`package.loaded["` + moduleName + `"] = nil`); err != nil {
			success = false
			logger.DebugCont("Failed to invalidate module %s: %v", moduleName, err)
		}
	}

	return success
}

// GetStateCount returns the number of initialized states
func (r *Runner) GetStateCount() int {
	r.mu.RLock()
	defer r.mu.RUnlock()

	count := 0
	for _, state := range r.states {
		if state != nil {
			count++
		}
	}

	return count
}

// GetActiveStateCount returns the number of states currently in use
func (r *Runner) GetActiveStateCount() int {
	r.mu.RLock()
	defer r.mu.RUnlock()

	count := 0
	for _, state := range r.states {
		if state != nil && state.inUse {
			count++
		}
	}

	return count
}