Moonshark/core/workers/eventloop_test.go

package workers

import (
	"context"
	"sync"
	"testing"
	"time"

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

// Helper function to create bytecode for testing
func createTestBytecode(t *testing.T, code string) []byte {
	state := luajit.New()
	if state == nil {
		t.Fatal("Failed to create Lua state")
	}
	defer state.Close()

	bytecode, err := state.CompileBytecode(code, "test")
	if err != nil {
		t.Fatalf("Failed to compile test bytecode: %v", err)
	}

	return bytecode
}

// Test creating a new event loop with default and custom configs
func TestNewEventLoop(t *testing.T) {
	tests := []struct {
		name        string
		config      EventLoopConfig
		expectError bool
	}{
		{
			name:        "Default config",
			config:      EventLoopConfig{},
			expectError: false,
		},
		{
			name: "Custom buffer size",
			config: EventLoopConfig{
				BufferSize: 200,
			},
			expectError: false,
		},
		{
			name: "Custom timeout",
			config: EventLoopConfig{
				Timeout: 5 * time.Second,
			},
			expectError: false,
		},
		{
			name: "With init function",
			config: EventLoopConfig{
				StateInit: func(state *luajit.State) error {
					return nil
				},
			},
			expectError: false,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			el, err := NewEventLoopWithConfig(tc.config)
			if tc.expectError {
				if err == nil {
					t.Errorf("Expected error but got nil")
				}
			} else {
				if err != nil {
					t.Errorf("Unexpected error: %v", err)
				}
				if el == nil {
					t.Errorf("Expected non-nil event loop")
				} else {
					el.Shutdown()
				}
			}
		})
	}
}

// Test basic job submission and execution
func TestEventLoopBasicSubmission(t *testing.T) {
	el, err := NewEventLoop()
	if err != nil {
		t.Fatalf("Failed to create event loop: %v", err)
	}
	defer el.Shutdown()

	// Simple return a value
	bytecode := createTestBytecode(t, "return 42")

	result, err := el.Submit(bytecode, nil)
	if err != nil {
		t.Fatalf("Failed to submit job: %v", err)
	}

	num, ok := result.(float64)
	if !ok {
		t.Fatalf("Expected float64 result, got %T", result)
	}

	if num != 42 {
		t.Errorf("Expected 42, got %f", num)
	}

	// Test more complex Lua code
	bytecode = createTestBytecode(t, `
		local result = 0
		for i = 1, 10 do
			result = result + i
		end
		return result
	`)

	result, err = el.Submit(bytecode, nil)
	if err != nil {
		t.Fatalf("Failed to submit job: %v", err)
	}

	num, ok = result.(float64)
	if !ok {
		t.Fatalf("Expected float64 result, got %T", result)
	}

	if num != 55 {
		t.Errorf("Expected 55, got %f", num)
	}
}

// Test context passing between Go and Lua
func TestEventLoopContext(t *testing.T) {
	el, err := NewEventLoop()
	if err != nil {
		t.Fatalf("Failed to create event loop: %v", err)
	}
	defer el.Shutdown()

	bytecode := createTestBytecode(t, `
		return {
			num = ctx.number,
			str = ctx.text,
			flag = ctx.enabled,
			list = {ctx.items[1], ctx.items[2], ctx.items[3]},
		}
	`)

	execCtx := NewContext()
	execCtx.Set("number", 42.5)
	execCtx.Set("text", "hello")
	execCtx.Set("enabled", true)
	execCtx.Set("items", []float64{10, 20, 30})

	result, err := el.Submit(bytecode, execCtx)
	if err != nil {
		t.Fatalf("Failed to submit job: %v", err)
	}

	// Result should be a map
	resultMap, ok := result.(map[string]any)
	if !ok {
		t.Fatalf("Expected map result, got %T", result)
	}

	// Check values
	if resultMap["num"] != 42.5 {
		t.Errorf("Expected num=42.5, got %v", resultMap["num"])
	}
	if resultMap["str"] != "hello" {
		t.Errorf("Expected str=hello, got %v", resultMap["str"])
	}
	if resultMap["flag"] != true {
		t.Errorf("Expected flag=true, got %v", resultMap["flag"])
	}

	arr, ok := resultMap["list"].([]float64)
	if !ok {
		t.Fatalf("Expected []float64, got %T", resultMap["list"])
	}

	expected := []float64{10, 20, 30}
	for i, v := range expected {
		if arr[i] != v {
			t.Errorf("Expected list[%d]=%f, got %f", i, v, arr[i])
		}
	}

	// Test complex nested context
	nestedCtx := NewContext()
	nestedCtx.Set("user", map[string]any{
		"id":   123,
		"name": "test user",
		"roles": []any{
			"admin",
			"editor",
		},
	})

	bytecode = createTestBytecode(t, `
		return {
			id = ctx.user.id,
			name = ctx.user.name,
			role1 = ctx.user.roles[1],
			role2 = ctx.user.roles[2],
		}
	`)

	result, err = el.Submit(bytecode, nestedCtx)
	if err != nil {
		t.Fatalf("Failed to submit job with nested context: %v", err)
	}

	resultMap, ok = result.(map[string]any)
	if !ok {
		t.Fatalf("Expected map result, got %T", result)
	}

	if resultMap["id"] != float64(123) {
		t.Errorf("Expected id=123, got %v", resultMap["id"])
	}
	if resultMap["name"] != "test user" {
		t.Errorf("Expected name='test user', got %v", resultMap["name"])
	}
	if resultMap["role1"] != "admin" {
		t.Errorf("Expected role1='admin', got %v", resultMap["role1"])
	}
	if resultMap["role2"] != "editor" {
		t.Errorf("Expected role2='editor', got %v", resultMap["role2"])
	}
}

// Test execution timeout
func TestEventLoopTimeout(t *testing.T) {
	// Create event loop with short timeout
	el, err := NewEventLoopWithConfig(EventLoopConfig{
		Timeout: 100 * time.Millisecond,
	})
	if err != nil {
		t.Fatalf("Failed to create event loop: %v", err)
	}
	defer el.Shutdown()

	// Create bytecode that runs for longer than the timeout
	bytecode := createTestBytecode(t, `
		-- Loop for 500ms
		local start = os.time()
		while os.difftime(os.time(), start) < 0.5 do end
		return "done"
	`)

	// This should time out
	_, err = el.Submit(bytecode, nil)
	if err != ErrExecutionTimeout {
		t.Errorf("Expected timeout error, got: %v", err)
	}

	// Now set a longer timeout and try again
	el.SetTimeout(1 * time.Second)

	// This should succeed
	result, err := el.Submit(bytecode, nil)
	if err != nil {
		t.Fatalf("Expected success with longer timeout, got: %v", err)
	}
	if result != "done" {
		t.Errorf("Expected 'done', got %v", result)
	}

	// Test per-call timeout with SubmitWithTimeout
	bytecode = createTestBytecode(t, `
		-- Loop for 300ms
		local start = os.time()
		while os.difftime(os.time(), start) < 0.3 do end
		return "done again"
	`)

	// This should time out with a custom timeout
	_, err = el.SubmitWithTimeout(bytecode, nil, 50*time.Millisecond)
	if err == nil {
		t.Errorf("Expected timeout error, got success")
	}

	// This should succeed with a longer custom timeout
	result, err = el.SubmitWithTimeout(bytecode, nil, 500*time.Millisecond)
	if err != nil {
		t.Fatalf("Expected success with custom timeout, got: %v", err)
	}
	if result != "done again" {
		t.Errorf("Expected 'done again', got %v", result)
	}
}

// Test module registration and execution
func TestEventLoopModules(t *testing.T) {
	// Define an init function that registers a simple "math" module
	mathInit := func(state *luajit.State) error {
		// Register the "add" function directly
		err := state.RegisterGoFunction("add", func(s *luajit.State) int {
			a := s.ToNumber(1)
			b := s.ToNumber(2)
			s.PushNumber(a + b)
			return 1 // Return one result
		})

		if err != nil {
			return err
		}

		// Register a math module with multiple functions
		mathFuncs := map[string]luajit.GoFunction{
			"multiply": func(s *luajit.State) int {
				a := s.ToNumber(1)
				b := s.ToNumber(2)
				s.PushNumber(a * b)
				return 1
			},
			"subtract": func(s *luajit.State) int {
				a := s.ToNumber(1)
				b := s.ToNumber(2)
				s.PushNumber(a - b)
				return 1
			},
		}

		return RegisterModule(state, "math2", mathFuncs)
	}

	// Create an event loop with our init function
	el, err := NewEventLoopWithInit(mathInit)
	if err != nil {
		t.Fatalf("Failed to create event loop: %v", err)
	}
	defer el.Shutdown()

	// Test the add function
	bytecode1 := createTestBytecode(t, "return add(5, 7)")
	result1, err := el.Submit(bytecode1, nil)
	if err != nil {
		t.Fatalf("Failed to call add function: %v", err)
	}

	num1, ok := result1.(float64)
	if !ok || num1 != 12 {
		t.Errorf("Expected add(5, 7) = 12, got %v", result1)
	}

	// Test the math2 module
	bytecode2 := createTestBytecode(t, "return math2.multiply(6, 8)")
	result2, err := el.Submit(bytecode2, nil)
	if err != nil {
		t.Fatalf("Failed to call math2.multiply: %v", err)
	}

	num2, ok := result2.(float64)
	if !ok || num2 != 48 {
		t.Errorf("Expected math2.multiply(6, 8) = 48, got %v", result2)
	}

	// Test multiple operations
	bytecode3 := createTestBytecode(t, `
		local a = add(10, 20)
		local b = math2.subtract(a, 5)
		return math2.multiply(b, 2)
	`)

	result3, err := el.Submit(bytecode3, nil)
	if err != nil {
		t.Fatalf("Failed to execute combined operations: %v", err)
	}

	num3, ok := result3.(float64)
	if !ok || num3 != 50 {
		t.Errorf("Expected ((10 + 20) - 5) * 2 = 50, got %v", result3)
	}
}

// Test combined module init functions
func TestEventLoopCombinedModules(t *testing.T) {
	// First init function adds a function to get a constant value
	init1 := func(state *luajit.State) error {
		return state.RegisterGoFunction("getAnswer", func(s *luajit.State) int {
			s.PushNumber(42)
			return 1
		})
	}

	// Second init function registers a function that multiplies a number by 2
	init2 := func(state *luajit.State) error {
		return state.RegisterGoFunction("double", func(s *luajit.State) int {
			n := s.ToNumber(1)
			s.PushNumber(n * 2)
			return 1
		})
	}

	// Combine the init functions
	combinedInit := CombineInitFuncs(init1, init2)

	// Create an event loop with the combined init function
	el, err := NewEventLoopWithInit(combinedInit)
	if err != nil {
		t.Fatalf("Failed to create event loop: %v", err)
	}
	defer el.Shutdown()

	// Test using both functions together in a single script
	bytecode := createTestBytecode(t, "return double(getAnswer())")
	result, err := el.Submit(bytecode, nil)
	if err != nil {
		t.Fatalf("Failed to execute: %v", err)
	}

	num, ok := result.(float64)
	if !ok || num != 84 {
		t.Errorf("Expected double(getAnswer()) = 84, got %v", result)
	}
}

// Test sandbox isolation between executions
func TestEventLoopSandboxIsolation(t *testing.T) {
	el, err := NewEventLoop()
	if err != nil {
		t.Fatalf("Failed to create event loop: %v", err)
	}
	defer el.Shutdown()

	// Create a script that tries to modify a global variable
	bytecode1 := createTestBytecode(t, `
		-- Set a "global" variable
		my_global = "test value"
		return true
	`)

	_, err = el.Submit(bytecode1, nil)
	if err != nil {
		t.Fatalf("Failed to execute first script: %v", err)
	}

	// Now try to access that variable from another script
	bytecode2 := createTestBytecode(t, `
		-- Try to access the previously set global
		return my_global ~= nil
	`)

	result, err := el.Submit(bytecode2, nil)
	if err != nil {
		t.Fatalf("Failed to execute second script: %v", err)
	}

	// The variable should not be accessible (sandbox isolation)
	if result.(bool) {
		t.Errorf("Expected sandbox isolation, but global variable was accessible")
	}
}

// Test error handling
func TestEventLoopErrorHandling(t *testing.T) {
	el, err := NewEventLoop()
	if err != nil {
		t.Fatalf("Failed to create event loop: %v", err)
	}
	defer el.Shutdown()

	// Test invalid bytecode
	_, err = el.Submit([]byte("not valid bytecode"), nil)
	if err == nil {
		t.Errorf("Expected error for invalid bytecode, got nil")
	}

	// Test Lua runtime error
	bytecode := createTestBytecode(t, `
		error("intentional error")
		return true
	`)

	_, err = el.Submit(bytecode, nil)
	if err == nil {
		t.Errorf("Expected error from Lua error() call, got nil")
	}

	// Test with nil context (should work fine)
	bytecode = createTestBytecode(t, "return ctx == nil")
	result, err := el.Submit(bytecode, nil)
	if err != nil {
		t.Errorf("Unexpected error with nil context: %v", err)
	}
	if result.(bool) != true {
		t.Errorf("Expected ctx to be nil in Lua, but it wasn't")
	}

	// Test access to restricted library
	bytecode = createTestBytecode(t, `
		-- Try to access io library directly
		return io ~= nil
	`)

	result, err = el.Submit(bytecode, nil)
	if err != nil {
		t.Fatalf("Failed to execute sandbox test: %v", err)
	}

	// io should not be directly accessible
	if result.(bool) {
		t.Errorf("Expected io library to be restricted, but it was accessible")
	}
}

// Test concurrent job submission
func TestEventLoopConcurrency(t *testing.T) {
	el, err := NewEventLoopWithConfig(EventLoopConfig{
		BufferSize: 100, // Buffer for concurrent submissions
		Timeout:    5 * time.Second,
	})
	if err != nil {
		t.Fatalf("Failed to create event loop: %v", err)
	}
	defer el.Shutdown()

	// Create bytecode that returns its input value
	bytecode := createTestBytecode(t, "return ctx.n")

	// Submit multiple jobs concurrently
	const jobCount = 50
	var wg sync.WaitGroup
	results := make([]int, jobCount)

	wg.Add(jobCount)
	for i := 0; i < jobCount; i++ {
		i := i // Capture loop variable
		go func() {
			defer wg.Done()

			// Create context with job number
			ctx := NewContext()
			ctx.Set("n", float64(i))

			// Submit job
			result, err := el.Submit(bytecode, ctx)
			if err != nil {
				t.Errorf("Job %d failed: %v", i, err)
				return
			}

			// Verify result matches job number
			num, ok := result.(float64)
			if !ok {
				t.Errorf("Job %d: expected float64, got %T", i, result)
				return
			}

			results[i] = int(num)
		}()
	}

	wg.Wait()

	// Verify all results
	for i, res := range results {
		if res != i && res != 0 { // 0 means error already logged
			t.Errorf("Expected result[%d] = %d, got %d", i, i, res)
		}
	}
}

// Test state consistency across multiple calls
func TestEventLoopStateConsistency(t *testing.T) {
	// Create an event loop with a module that maintains count between calls
	initFunc := func(state *luajit.State) error {
		// Create a closure that increments a counter in upvalue
		code := `
			-- Create a counter with initial value 0
			local counter = 0

			-- Create a function that returns and increments the counter
			function get_next_count()
				local current = counter
				counter = counter + 1
				return current
			end
		`
		return state.DoString(code)
	}

	el, err := NewEventLoopWithInit(initFunc)
	if err != nil {
		t.Fatalf("Failed to create event loop: %v", err)
	}
	defer el.Shutdown()

	// Now run multiple scripts that call the counter function
	bytecode := createTestBytecode(t, "return get_next_count()")

	// Each call should return an incremented value
	for i := 0; i < 5; i++ {
		result, err := el.Submit(bytecode, nil)
		if err != nil {
			t.Fatalf("Call %d failed: %v", i, err)
		}

		num, ok := result.(float64)
		if !ok {
			t.Fatalf("Expected float64 result, got %T", result)
		}

		if int(num) != i {
			t.Errorf("Expected count %d, got %d", i, int(num))
		}
	}
}

// Test shutdown and cleanup
func TestEventLoopShutdown(t *testing.T) {
	el, err := NewEventLoop()
	if err != nil {
		t.Fatalf("Failed to create event loop: %v", err)
	}

	// Submit a job to verify it works
	bytecode := createTestBytecode(t, "return 42")
	_, err = el.Submit(bytecode, nil)
	if err != nil {
		t.Fatalf("Failed to submit job: %v", err)
	}

	// Shutdown
	if err := el.Shutdown(); err != nil {
		t.Errorf("Shutdown failed: %v", err)
	}

	// Submit after shutdown should fail
	_, err = el.Submit(bytecode, nil)
	if err != ErrLoopClosed {
		t.Errorf("Expected ErrLoopClosed, got %v", err)
	}

	// Second shutdown should return error
	if err := el.Shutdown(); err != ErrLoopClosed {
		t.Errorf("Expected ErrLoopClosed on second shutdown, got %v", err)
	}
}

// Test high load with multiple sequential and concurrent jobs
func TestEventLoopHighLoad(t *testing.T) {
	el, err := NewEventLoopWithConfig(EventLoopConfig{
		BufferSize: 1000, // Large buffer for high load
		Timeout:    5 * time.Second,
	})
	if err != nil {
		t.Fatalf("Failed to create event loop: %v", err)
	}
	defer el.Shutdown()

	// Sequential load test
	bytecode := createTestBytecode(t, `
		-- Do some work
		local result = 0
		for i = 1, 1000 do
			result = result + i
		end
		return result
	`)

	start := time.Now()
	for i := 0; i < 100; i++ {
		_, err := el.Submit(bytecode, nil)
		if err != nil {
			t.Fatalf("Sequential job %d failed: %v", i, err)
		}
	}
	seqDuration := time.Since(start)
	t.Logf("Sequential load test: 100 jobs in %v", seqDuration)

	// Concurrent load test
	start = time.Now()
	var wg sync.WaitGroup
	wg.Add(100)
	for i := 0; i < 100; i++ {
		go func() {
			defer wg.Done()
			_, err := el.Submit(bytecode, nil)
			if err != nil {
				t.Errorf("Concurrent job failed: %v", err)
			}
		}()
	}
	wg.Wait()
	concDuration := time.Since(start)
	t.Logf("Concurrent load test: 100 jobs in %v", concDuration)
}

// Test context cancellation
func TestEventLoopCancel(t *testing.T) {
	el, err := NewEventLoop()
	if err != nil {
		t.Fatalf("Failed to create event loop: %v", err)
	}
	defer el.Shutdown()

	// Create a long-running script
	bytecode := createTestBytecode(t, `
		-- Sleep for 500ms
		local start = os.time()
		while os.difftime(os.time(), start) < 0.5 do end
		return "done"
	`)

	// Create a context that we can cancel
	ctx, cancel := context.WithCancel(context.Background())

	// Start execution in a goroutine
	resultCh := make(chan any, 1)
	errCh := make(chan error, 1)
	go func() {
		res, err := el.SubmitWithContext(ctx, bytecode, nil)
		if err != nil {
			errCh <- err
		} else {
			resultCh <- res
		}
	}()

	// Cancel quickly
	time.Sleep(50 * time.Millisecond)
	cancel()

	// Should get cancellation error
	select {
	case err := <-errCh:
		if ctx.Err() == nil || err == nil {
			t.Errorf("Expected context cancellation error")
		}
	case res := <-resultCh:
		t.Errorf("Expected cancellation, got result: %v", res)
	case <-time.After(1 * time.Second):
		t.Errorf("Timed out waiting for cancellation")
	}
}