package runner

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"

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

// Common errors
var (
	ErrRunnerClosed = errors.New("lua runner is closed")
	ErrInitFailed   = errors.New("initialization failed")
)

// StateInitFunc is a function that initializes a Lua state
type StateInitFunc func(*luajit.State) error

// LuaRunner runs Lua scripts using a single Lua state
type LuaRunner struct {
	state      *luajit.State  // The Lua state
	jobQueue   chan job       // Channel for incoming jobs
	isRunning  atomic.Bool    // Flag indicating if the runner is active
	mu         sync.RWMutex   // Mutex for thread safety
	wg         sync.WaitGroup // WaitGroup for clean shutdown
	initFunc   StateInitFunc  // Optional function to initialize Lua state
	bufferSize int            // Size of the job queue buffer
}

// NewRunner creates a new LuaRunner
func NewRunner(options ...RunnerOption) (*LuaRunner, error) {
	// Default configuration
	runner := &LuaRunner{
		bufferSize: 10, // Default buffer size
	}

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

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

	// Create job queue
	runner.jobQueue = make(chan job, runner.bufferSize)
	runner.isRunning.Store(true)

	// Set up sandbox
	if err := runner.setupSandbox(); err != nil {
		state.Close()
		return nil, ErrInitFailed
	}

	// Run init function if provided
	if runner.initFunc != nil {
		if err := runner.initFunc(state); err != nil {
			state.Close()
			return nil, ErrInitFailed
		}
	}

	// Start the event loop
	runner.wg.Add(1)
	go runner.eventLoop()

	return runner, nil
}

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

// WithBufferSize sets the job queue buffer size
func WithBufferSize(size int) RunnerOption {
	return func(r *LuaRunner) {
		if size > 0 {
			r.bufferSize = size
		}
	}
}

// WithInitFunc sets the init function for the Lua state
func WithInitFunc(initFunc StateInitFunc) RunnerOption {
	return func(r *LuaRunner) {
		r.initFunc = initFunc
	}
}

// setupSandbox initializes the sandbox environment
func (r *LuaRunner) setupSandbox() error {
	// This is the Lua script that creates our sandbox function
	setupScript := `
	-- Create a function to run code in a sandbox environment
	function __create_sandbox()
		-- Create new environment table
		local env = {}
		
		-- Add standard library modules (can be restricted as needed)
		env.string = string
		env.table = table
		env.math = math
		env.os = { 
			time = os.time, 
			date = os.date,
			difftime = os.difftime,
			clock = os.clock
		}
		env.tonumber = tonumber
		env.tostring = tostring
		env.type = type
		env.pairs = pairs
		env.ipairs = ipairs
		env.next = next
		env.select = select
		env.unpack = unpack
		env.pcall = pcall
		env.xpcall = xpcall
		env.error = error
		env.assert = assert
		
		-- Allow access to package.loaded for modules
		env.require = function(name)
			return package.loaded[name]
		end
		
		-- Create metatable to restrict access to _G
		local mt = {
			__index = function(t, k)
				-- First check in env table
				local v = rawget(env, k)
				if v ~= nil then return v end
				
				-- If not found, check for registered modules/functions
				local moduleValue = _G[k]
				if type(moduleValue) == "table" or 
				   type(moduleValue) == "function" then
					return moduleValue
				end
				
				return nil
			end,
			__newindex = function(t, k, v)
				rawset(env, k, v)
			end
		}
		
		setmetatable(env, mt)
		return env
	end
	
	-- Create function to execute code with a sandbox
	function __run_sandboxed(f, ctx)
		local env = __create_sandbox()
		
		-- Add context to the environment if provided
		if ctx then
			env.ctx = ctx
		end
		
		-- Set the environment and run the function
		setfenv(f, env)
		return f()
	end
	`

	return r.state.DoString(setupScript)
}

// eventLoop processes jobs from the queue
func (r *LuaRunner) eventLoop() {
	defer r.wg.Done()
	defer r.state.Close()

	// Process jobs until closure
	for job := range r.jobQueue {
		// Execute the job and send result
		result := r.executeJob(job)
		select {
		case job.Result <- result:
			// Result sent successfully
		default:
			// Result channel closed or full, discard the result
		}
	}
}

// executeJob runs a script in the sandbox environment
func (r *LuaRunner) executeJob(j job) JobResult {
	// Re-run init function if needed
	if r.initFunc != nil {
		if err := r.initFunc(r.state); err != nil {
			return JobResult{nil, err}
		}
	}

	// Set up context if provided
	if j.Context != nil {
		// Push context table
		r.state.NewTable()

		// Add values to context table
		for key, value := range j.Context.Values {
			// Push key
			r.state.PushString(key)

			// Push value
			if err := r.state.PushValue(value); err != nil {
				return JobResult{nil, err}
			}

			// Set table[key] = value
			r.state.SetTable(-3)
		}
	} else {
		// Push nil if no context
		r.state.PushNil()
	}

	// Load bytecode
	if err := r.state.LoadBytecode(j.Bytecode, "script"); err != nil {
		r.state.Pop(1) // Pop context
		return JobResult{nil, err}
	}

	// Get the sandbox runner function
	r.state.GetGlobal("__run_sandboxed")

	// Push loaded function and context as arguments
	r.state.PushCopy(-2) // Copy the loaded function
	r.state.PushCopy(-4) // Copy the context table or nil

	// Remove the original function and context
	r.state.Remove(-5) // Remove original context
	r.state.Remove(-4) // Remove original function

	// Call the sandbox runner with 2 args (function and context), expecting 1 result
	if err := r.state.Call(2, 1); err != nil {
		return JobResult{nil, err}
	}

	// Get result
	value, err := r.state.ToValue(-1)
	r.state.Pop(1) // Pop result

	return JobResult{value, err}
}

// RunWithContext executes a Lua script with context and timeout
func (r *LuaRunner) RunWithContext(ctx context.Context, bytecode []byte, execCtx *Context) (any, error) {
	r.mu.RLock()
	if !r.isRunning.Load() {
		r.mu.RUnlock()
		return nil, ErrRunnerClosed
	}
	r.mu.RUnlock()

	resultChan := make(chan JobResult, 1)
	j := job{
		Bytecode: bytecode,
		Context:  execCtx,
		Result:   resultChan,
	}

	// Submit job with context
	select {
	case r.jobQueue <- j:
		// Job submitted
	case <-ctx.Done():
		return nil, ctx.Err()
	}

	// Wait for result with context
	select {
	case result := <-resultChan:
		return result.Value, result.Error
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

// Run executes a Lua script
func (r *LuaRunner) Run(bytecode []byte, execCtx *Context) (any, error) {
	return r.RunWithContext(context.Background(), bytecode, execCtx)
}

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

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

	r.isRunning.Store(false)
	close(r.jobQueue)

	// Wait for event loop to finish
	r.wg.Wait()

	return nil
}