--[[
The purpose of this script is to make it easier to call lua code in the
Cheat Engine process from the game.  The following commands will be made
available:

CELUA_SYNC(functionName, arg1, arg2, ...)
CELUA_ASYNC(functionName, arg1, arg2, ...)

Reference: https://forum.cheatengine.org/viewtopic.php?t=605733

These should work in 32-bit and 64-bit tables.  All flags and registers should
be preserved except for eax/rax which will hold the return value of the
function.  There are a few caveats with the values you can use:

32-bit - the string you pass for each argument is used as the argument for
  the push instruction.  So if you call this:
      CELUA_ASYNC(MyFunction, eax, [ebp+08], [pPlayer])
  these push instructions will be used (args are pushed in reverse order)
      push [pPlayer]
      push [ebp+08]
      push eax
  So you can't use things like XMM registers, 16-bit values, etc.  You'll have
  to save those to memory or the stack yourself and load a register with the
  address or something.  Also since values are pushed onto the stack, a
  stack reference like '[esp+28]' would not be what you expect.

64-bit - the string you pass for each argument is loaded into rax and pushed
  onto the stack for sending to lua.  Since the arguments are 64-bit, you should
  load any 32-bit values into 32 bit registers and pass them, or put them
  on the stack or in memory and pass a pointer to them.  For instance:
      CELUA_ASYNC(MyFunction, rcx, [ebp+10], [globals])
  Will use these instructions:
      mov rax,[globals]
      push rax
      mov rax,[ebp+10]
      push rax
      mov rax,rcx
      push rax
  In addition since rax is used you cannot use rax for any but the last argument
  or it would be overwritten by earlier instructions.  Finally you can't load
  far addresses directly in x64 so if the symbol globals isn't in the same
  memory space, you would have to load a register with the the address, like:
      mov rdx,globals
      CELUA_ASYNC(MyFunction, [rdx], [rax+8]) // rax as final argument is fine

--]]

CELUAOverrides = {}
local me = CELUAOverrides

-- split() helper function
function me.split(s, pattern, maxsplit)
  local pattern = pattern or ' '
  local maxsplit = maxsplit or -1
  local s = s
  local t = {}
  local patsz = #pattern
  while maxsplit ~= 0 do
    local curpos = 1
    local found = string.find(s, pattern)
    if found ~= nil then
      table.insert(t, string.sub(s, curpos, found - 1))
      curpos = found + patsz
      s = string.sub(s, curpos)
    else
      table.insert(t, string.sub(s, curpos))
      break
    end
    maxsplit = maxsplit - 1
    if maxsplit == 0 then
      table.insert(t, string.sub(s, curpos - patsz - 1))
    end
  end
  return t
end

-- trim() helper function
function me.trim(s)
  return s:gsub("^%s+", ""):gsub("%s+$", "")
end

local aaInitialize_i386 = [[
loadlibrary(luaclient-i386.dll)
globalalloc(__celua_initialized__,$1000)
__celua_initialized__:
  dd 0               // result will go here, should be FFFFFFFF on success
  dd 1234567f        // magic number
  db 'CELUASERVER',0 // server name

__celua_initialized__+100:
  // code
  push __celua_initialized__+8    // server name
  call CELUA_Initialize
  mov [__celua_initialized__],eax // save result
  ret

CREATETHREAD(__celua_initialized__+100)
]]

local aaInitialize_x86_64 = [[
globalalloc(__celua_initialized__,$1000)

loadlibrary(luaclient-x86_64.dll)

__celua_initialized__:
  dd 0               // result will go here, should be FFFFFFFF on success
  dd 1234567f        // magic number
  db 'CELUASERVER',0 // server name

__celua_initialized__+100:
  mov rcx,__celua_initialized__+8    // server name
  call CELUA_Initialize
  mov [__celua_initialized__],eax // save result
  ret

CREATETHREAD(__celua_initialized__+100)
]]

local aaFunction_i386 = [[
globalalloc(__celua_FUNCTION_MYFUNCTIONNAME__,$1000)

__celua_FUNCTION_MYFUNCTIONNAME__:
  dd 0               // result will go here, should be non-zero on success
  dd 1234567f        // magic number
  db 'MYFUNCTIONNAME',0  // function name

__celua_FUNCTION_MYFUNCTIONNAME__+100:
  // code

  mov eax, [__celua_FUNCTION_MYFUNCTIONNAME__]
  cmp eax,0
  jne @F
  cmp [__celua_initialized__],ffffffff
  jne @F
  push __celua_FUNCTION_MYFUNCTIONNAME__+8    // function name
  call CELUA_GetFunctionReferenceFromName
  mov [__celua_FUNCTION_MYFUNCTIONNAME__],eax // save result
@@:
  ret
]]

local aaFunction_x86_64 = [[
globalalloc(__celua_FUNCTION_MYFUNCTIONNAME__,$1000)

__celua_FUNCTION_MYFUNCTIONNAME__:
  dd 0               // result will go here, should be non-zero on success
  dd 1234567f        // magic number
  db 'MyFunction',0  // function name

__celua_FUNCTION_MYFUNCTIONNAME__+100:
  mov eax, [__celua_FUNCTION_MYFUNCTIONNAME__]
  test eax,eax
  jnz @F
  cmp [__celua_initialized__],ffffffff
  jne @F
  mov rcx, __celua_FUNCTION_MYFUNCTIONNAME__+8    // function name
  call CELUA_GetFunctionReferenceFromName
  mov [__celua_FUNCTION_MYFUNCTIONNAME__],eax // save result
@@:
  ret
]]


-- Ensure that __celua_initialized__ is setup correctly
-- returns true if it's ok, or false,errorMessage if not
function me.EnsureInitialized()
  local symbolname = '__celua_initialized__'

  -- see if the symbol is already defined
  local addr = getAddressSafe(symbolname)
  if addr ~= nil then
    -- we found it, ensure that it contains correct data
    local magicNumber = readInteger(symbolname..'+4')
    if magicNumber == 0x1234567f then return true end -- already created
    unregisterSymbol(symbolname) -- invalid data at that address, unregister
  end

  -- assemble appropriate script
  if targetIs64Bit() then
    if not autoAssemble(aaInitialize_x86_64) then return false, 'Error assembling initialize script' end
  else
    if not autoAssemble(aaInitialize_i386) then return false, 'Error assembling initialize script' end
  end
  return true
end

function me.GetFunctionSymbol(functionName)
  if type(functionName) ~= 'string' then return nil end
  if string.len(functionName) <= 0 then return nil end
  return '__celua_FUNCTION_'..functionName..'__'
end

function me.EnsureFunction(functionName)
  local symbolname = me.GetFunctionSymbol(functionName)
  if symbolname == nil then return false, 'Unable to create symbol for function name "'..tostring(functionName)..'"' end
  local addr = getAddressSafe(symbolname)
  if addr ~= nil then
    -- we found it, ensure that it contains correct data
    local magicNumber = readInteger(symbolname..'+4')
    if magicNumber == 0x1234567f then return true end -- already created
    unregisterSymbol(symbolname) -- invalid data at that address, unregister
  end
  local aa
  if targetIs64Bit() then
    aa = string.gsub(aaFunction_x86_64, 'MYFUNCTIONNAME', functionName)
  else
    aa = string.gsub(aaFunction_i386, 'MYFUNCTIONNAME', functionName)
  end
  if not autoAssemble(aa) then
    print("Error assembling lua function aa code:")
    print(string.gsub(aa, '\n', '\r\n'))
    return false, 'Error assembling function script'
  end
  return true
end

local labelCounter = math.random(0x70000000) + 0x10000000

local aaPerformCall_i386 = [[
label(${LABEL_NAME})
pushad
pushfd

${PARAMETERS}

call __celua_FUNCTION_${FUNCTION_NAME}__+100
test eax,eax
jz ${LABEL_NAME}

mov ecx, esp
push ${ASYNC_VALUE}
push ecx
push ${PARAMETERS_COUNT}
push eax
call CELUA_ExecuteFunctionByReference

${LABEL_NAME}:
add esp,${PARAMETERS_SIZE}
mov [esp+20],eax

popfd
popad
]]

local aaPerformCall_x86_64 = [[
label(${LABEL_NAME})
pushfq
push rcx
push rdx
push r8
push r9
push r10
push r11
push rsi

${PARAMETERS}

mov rsi, rsp
and rsp,fffffffffffffff0

call __celua_FUNCTION_${FUNCTION_NAME}__+100
test eax,eax
jz ${LABEL_NAME}

mov r8, rsi
mov r9, ${ASYNC_VALUE}
mov edx, ${PARAMETERS_COUNT}
mov ecx, eax
call CELUA_ExecuteFunctionByReference

${LABEL_NAME}:
mov rsp, rsi
add esp, ${PARAMETERS_SIZE}
pop rsi
pop r11
pop r10
pop r9
pop r8
pop rdx
pop rcx
popfq
]]

function me.AACall(functionName, parameters, async)
  local label = 'skip_'..functionName..'_'..string.format('%2x', labelCounter)
  labelCounter = labelCounter + 1

  local pushes ={}

  local asyncValueString = '0'
  if async then asyncValueString = '1' end

  if targetIs64Bit() then
    for i=#parameters,1,-1 do
      table.insert(pushes, 'mov rax, '..parameters[i])
      table.insert(pushes, 'push rax')
    end

    local result = aaPerformCall_x86_64
    result = result:gsub('${LABEL_NAME}', label)
    result = result:gsub('${FUNCTION_NAME}', functionName)
    result = result:gsub('${PARAMETERS}', table.concat(pushes, '\n'))
    result = result:gsub('${PARAMETERS_SIZE}', string.format('%x', (#parameters)*8))
    result = result:gsub('${PARAMETERS_COUNT}', string.format('%x', #parameters))
    result = result:gsub('${ASYNC_VALUE}', asyncValueString)

    return result
  else
    for i=#parameters,1,-1 do
      table.insert(pushes, 'push '..parameters[i])
    end

    local result = aaPerformCall_i386
    result = result:gsub('${LABEL_NAME}', label)
    result = result:gsub('${FUNCTION_NAME}', functionName)
    result = result:gsub('${PARAMETERS}', table.concat(pushes, '\n'))
    result = result:gsub('${PARAMETERS_SIZE}', string.format('%x', (#parameters)*4))
    result = result:gsub('${PARAMETERS_COUNT}', string.format('%x', #parameters))
    result = result:gsub('${ASYNC_VALUE}', asyncValueString)
    return result
  end
end

function me.AAOverride(parameters, syntaxCheckOnly, async)
  -- open the server ourselves
  openLuaServer('CELUASERVER')

  local result, msg = me.EnsureInitialized()
  if not result then return nil, msg end

  local vals = me.split(parameters, ',')
  local functionName = me.trim(vals[1])
  result, msg = me.EnsureFunction(functionName)
  if not result then return nil, msg end

  -- now we have to create the actual script
  local actualParameters = {}
  for i = 2,#vals do
    local value = me.trim(vals[i])
    table.insert(actualParameters, value)
  end
  local final = me.AACall(functionName, actualParameters, async)

  --print(final)
  --showMessage(final)
  return final
end

-- sync version
function me.AAOverride_SYNC(parameters, syntaxCheckOnly)
  return me.AAOverride(parameters, syntaxCheckOnly, false)
end

-- async version
function me.AAOverride_ASYNC(parameters, syntaxCheckOnly)
  return me.AAOverride(parameters, syntaxCheckOnly, true)
end

-- register the actual commands
registerAutoAssemblerCommand('CELUA_SYNC', CELUAOverrides.AAOverride_SYNC)
registerAutoAssemblerCommand('CELUA_ASYNC', CELUAOverrides.AAOverride_ASYNC)

--[[

return CELUAOverrides.EnsureInitialized()
return CELUAOverrides.AAOverride('MyFunction, eax, esi, [ebp+08]')

--]]