package runner

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

// Sandbox manages a sandboxed Lua environment
type Sandbox struct {
	modules     map[string]any // Custom modules for environment
	initialized bool           // Whether base environment is initialized
}

// NewSandbox creates a new sandbox
func NewSandbox() *Sandbox {
	s := &Sandbox{
		modules:     make(map[string]any),
		initialized: false,
	}

	return s
}

// AddModule adds a module to the sandbox environment
func (s *Sandbox) AddModule(name string, module any) {
	s.modules[name] = module
}

// Setup initializes the sandbox in a Lua state
func (s *Sandbox) Setup(state *luajit.State) error {
	// Register modules
	if err := s.registerModules(state); err != nil {
		return err
	}

	// Create high-performance persistent environment
	return state.DoString(`
    -- Global shared environment (created once)
    __env_system = __env_system or {
        base_env = nil,    -- Template environment
        initialized = false, -- Initialization flag
        env_pool = {},     -- Pre-allocated environment pool
        pool_size = 0,     -- Current pool size
        max_pool_size = 8  -- Maximum pool size
    }

    -- Initialize base environment once
    if not __env_system.initialized then
        -- Create base environment with all standard libraries
        local base = {}

        -- Safe standard libraries
        base.string = string
        base.table = table
        base.math = math
        base.os = {
            time = os.time,
            date = os.date,
            difftime = os.difftime,
            clock = os.clock
        }

        -- Basic functions
        base.print = print
        base.tonumber = tonumber
        base.tostring = tostring
        base.type = type
        base.pairs = pairs
        base.ipairs = ipairs
        base.next = next
        base.select = select
        base.unpack = unpack
        base.pcall = pcall
        base.xpcall = xpcall
        base.error = error
        base.assert = assert

        -- Package system is shared for performance
        base.package = {
            loaded = package.loaded,
            path = package.path,
            preload = package.preload
        }

        -- Add HTTP module explicitly to the base environment
        base.http = http

        -- Add registered custom modules
        if __sandbox_modules then
            for name, mod in pairs(__sandbox_modules) do
                base[name] = mod
            end
        end

        -- Store base environment
        __env_system.base_env = base
        __env_system.initialized = true
    end

		-- Global variable for tracking current environment
		__last_env = nil

		-- Fast environment creation with pre-allocation
		function __get_sandbox_env(ctx)
			local env

			-- Try to reuse from pool
			if __env_system.pool_size > 0 then
				env = table.remove(__env_system.env_pool)
				__env_system.pool_size = __env_system.pool_size - 1

				-- Clear any previous context
				env.ctx = ctx or nil
				-- Clear any previous response
				env._response = nil
			else
				-- Create new environment with metatable inheritance
				env = setmetatable({}, {
					__index = __env_system.base_env
				})

				-- Set context if provided
				if ctx then
					env.ctx = ctx
				end

				-- Install the fast require implementation
				env.require = function(modname)
					return __fast_require(env, modname)
				end
			end

			-- Store reference to current environment
			__last_env = env

			return env
		end

		-- Return environment to pool for reuse
		function __recycle_env(env)
			-- Only recycle if pool isn't full
			if __env_system.pool_size < __env_system.max_pool_size then
				-- Clear context reference to avoid memory leaks
				env.ctx = nil
				-- Don't clear response data - we need it for extraction

				-- Add to pool
				table.insert(__env_system.env_pool, env)
				__env_system.pool_size = __env_system.pool_size + 1
			end
		end

		-- Hyper-optimized sandbox executor
		function __execute_sandbox(bytecode, ctx)
			-- Get environment (from pool if available)
			local env = __get_sandbox_env(ctx)

			-- Set environment for bytecode
			setfenv(bytecode, env)

			-- Execute with protected call
			local success, result = pcall(bytecode)

			-- Recycle environment for future use
			__recycle_env(env)

			-- Process result
			if not success then
				error(result, 0)
			end

			return result
		end

		-- Run minimal GC for overall health
		collectgarbage("step", 10)
	`)
}

// registerModules registers custom modules in the Lua state
func (s *Sandbox) registerModules(state *luajit.State) error {
	// Create or get module registry table
	state.GetGlobal("__sandbox_modules")
	if state.IsNil(-1) {
		// Table doesn't exist, create it
		state.Pop(1)
		state.NewTable()
		state.SetGlobal("__sandbox_modules")
		state.GetGlobal("__sandbox_modules")
	}

	// Add modules to registry
	for name, module := range s.modules {
		state.PushString(name)
		if err := state.PushValue(module); err != nil {
			state.Pop(2)
			return err
		}
		state.SetTable(-3)
	}

	// Pop module table
	state.Pop(1)
	return nil
}

// Execute runs bytecode in the sandbox
func (s *Sandbox) Execute(state *luajit.State, bytecode []byte, ctx map[string]any) (any, error) {
	// Update custom modules if needed
	if !s.initialized {
		if err := s.registerModules(state); err != nil {
			return nil, err
		}
		s.initialized = true
	}

	// Load bytecode
	if err := state.LoadBytecode(bytecode, "script"); err != nil {
		return nil, err
	}

	// Create context table if provided
	if len(ctx) > 0 {
		// Preallocate table with appropriate size
		state.CreateTable(0, len(ctx))

		// Add context entries
		for k, v := range ctx {
			state.PushString(k)
			if err := state.PushValue(v); err != nil {
				state.Pop(2)
				return nil, err
			}
			state.SetTable(-3)
		}
	} else {
		state.PushNil() // No context
	}

	// Get optimized sandbox executor
	state.GetGlobal("__execute_sandbox")

	// Setup call with correct argument order
	state.PushCopy(-3) // Copy bytecode function
	state.PushCopy(-3) // Copy context

	// Clean up stack
	state.Remove(-5) // Remove original bytecode
	state.Remove(-4) // Remove original context

	// Call optimized sandbox executor
	if err := state.Call(2, 1); err != nil {
		return nil, err
	}

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

	// Check if HTTP response was set
	httpResponse, hasHTTPResponse := GetHTTPResponse(state)
	if hasHTTPResponse {
		httpResponse.Body = result
		return httpResponse, err
	}

	return result, err
}