package luajit

/*
#include <lua.h>
#include <lauxlib.h>
#include <stdlib.h>

extern int goFunctionWrapper(lua_State* L);

static int get_upvalue_index(int i) {
    return -10002 - i;  // LUA_GLOBALSINDEX - i
}
*/
import "C"
import (
	"fmt"
	"sync"
	"unsafe"
)

type GoFunction func(*State) int

var (
	functionRegistry = struct {
		sync.RWMutex
		funcs map[unsafe.Pointer]GoFunction
	}{
		funcs: make(map[unsafe.Pointer]GoFunction),
	}
)

//export goFunctionWrapper
func goFunctionWrapper(L *C.lua_State) C.int {
	state := &State{L: L, safeStack: true}

	// Get upvalue using standard Lua 5.1 macro
	ptr := C.lua_touserdata(L, C.get_upvalue_index(1))
	if ptr == nil {
		state.PushString("error: function not found")
		return -1
	}

	functionRegistry.RLock()
	fn, ok := functionRegistry.funcs[ptr]
	functionRegistry.RUnlock()

	if !ok {
		state.PushString("error: function not found in registry")
		return -1
	}

	result := fn(state)
	return C.int(result)
}

func (s *State) PushGoFunction(fn GoFunction) error {
	// Push lightuserdata as upvalue and create closure
	ptr := C.malloc(1)
	if ptr == nil {
		return fmt.Errorf("failed to allocate memory for function pointer")
	}

	functionRegistry.Lock()
	functionRegistry.funcs[ptr] = fn
	functionRegistry.Unlock()

	C.lua_pushlightuserdata(s.L, ptr)
	C.lua_pushcclosure(s.L, (*[0]byte)(C.goFunctionWrapper), 1)
	return nil
}

func (s *State) RegisterGoFunction(name string, fn GoFunction) error {
	if err := s.PushGoFunction(fn); err != nil {
		return err
	}

	cname := C.CString(name)
	defer C.free(unsafe.Pointer(cname))
	C.lua_setfield(s.L, C.LUA_GLOBALSINDEX, cname)
	return nil
}

func (s *State) UnregisterGoFunction(name string) {
	s.PushNil()
	cname := C.CString(name)
	defer C.free(unsafe.Pointer(cname))
	C.lua_setfield(s.L, C.LUA_GLOBALSINDEX, cname)
}

func (s *State) Cleanup() {
	functionRegistry.Lock()
	defer functionRegistry.Unlock()

	for ptr := range functionRegistry.funcs {
		C.free(ptr)
	}
	functionRegistry.funcs = make(map[unsafe.Pointer]GoFunction)
}