local NAME = 'I2CEPlugins.AutoAssemblerTemplates'
local CLASS_NAME = '[AutoAssemblerTemplates]'
local VERSION = '1.0.1'
local DESCRIPTION = 'I2 Cheat Engine Lua Modules'
local AUTHOR = 'TheyCallMeTim13@Gmail.com'
if BEBUG then print(NAME, ...) end

--[[
Template Variable:
	Address : number
		The address of the injection point.
	AddressString : string
		The address of the injection point as a string. It will be in "module+offset" or hex format.
	AobAddress : number
		The address of the AOB start (first selected line in the memory view form).
	AobAddressString : string
		Same as address but for "AobAddress".
	OriginalBytes : table/array : number
		A table of numbers (base 10) of the original bytes.
	OriginalBytesString : string
		A string of the original bytes in a "XX XX" (base 16) format.
	ModuleName : string
		A string of the module name if the injection point and full AOB is in a module.
	Signature : string
		The AOB in a signature format, for now it just removes hard coded address and 4 byte offsets.
	AOB : table/array : number
		A table of numbers (base 10) of the raw AOB.
	AobOffset : number
		The difference between "AobAddress" and "Address" or nil if they are the same.
	AobOffsetString : string
		The AOB offset in hex format with the appropriate sign (e.g.: "+10" or "-10").
		Or nil if there is no aob offset.
	PointerDefault : string
		Either "dd 0" or "dq 0" based on if the target is 64 bit.
	BaseAddressRegistry : string
		If opcode in the format of "[reg+n]" is found this would be set to "reg".
	BaseAddressOffset : string
		
	OriginalOpcode : table/array : string
		A table of strings of the original opcode lines.
	InjectionInfo : table/array : string
		Disassembler lines before and after the injection point.
	NopBytes : table
		A table of numbers (base 10) of the nop bytes (i.e.: 0x90) for the injection point 
		if it's not equal to 5 bytes. It will be an empty table if no nops are needed.
	NopBytesString : string
		The string of the nop bytes in a "XX XX" (base 16) format.
	CEVersion : number
		The current CE version.
	Date : string
		The current date.
	HookName : string
		If the template settings has "AskForHookName" set to true it will prompt for a hook name.
		Then title case the string (e.g.: "some test hook" will be come "Some Test Hook"), 
		and set this to that.
	HookNameParsed : string
		If the template settings has "AskForHookName" set to true it will prompt for a hook name.
		Then title case and remove spaces in the string 
		(e.g.: "some test hook" will be come "SomeTestHook"), and set this to that.
	FinalCompilation : boolean
		Only for the settings files as they get loaded once when CE loads then again at the start of 
		template rendering to get the template settings, and again just before rendering the 
		template but only on the last load will the environment be setup for the setting file. 
		This is used to know if the environment is setup if more environment setup is done in 
		the setting file.
		
Template Settings:
	Caption : string
		The caption of the template's menu item.
	Shortcut : string
		The shortcut of the template's menu item.
	AskForInjectionAddress : boolean
		Will prompt for a different address than the selected memory view lines if set to true.
		Allows the injection point to be offset from the AOB.
	AskForHookName : boolean
		Will prompt for a hook name if set to true.
	InjectionInfoLinesCount : number
		Will override the default number of lines in the "InjectionInfo" table.
	AppendToHookName : string
		Will be appended to the hook name if not already (if "AskForHookName" is true).
	WildCard : string/char
		Will override the default wild card if set.
]]

pcall(require, 'I2CETLua.addressList')
local status, _ = pcall(require, 'I2CETLua.TemplateEngine')
if not status then
	status, _ = pcall(require, 'TemplateEngine')
end
if not status then
	status, _ = pcall(dofile, getAutorunPath()..'TemplateEngine.lua')
end
if not status then
	print('"TemplateEngine" not found, it is required for "'..NAME..'" to work.')
	return
end
local sep = package.config:sub(1,1)


local WildCard = 'x'
local TemplatesFolder = getAutorunPath()..'Templates'..sep..'cea'
local CeaFileExtension = 'CEA'

if os.getenv('CE_LUA_PROJECTS_FOLDER') then
	TemplatesFolder = os.getenv('CE_LUA_PROJECTS_FOLDER')..sep..'Templates'..sep..'cea'
end


local function exists(name)
	if not name then return end
	return os.rename(name, name) and true or false
end

local function generateSignature(wildCard)
	wildCard = wildCard or WildCard
	local disassemblerView = getVisibleDisassembler()
	local dv_address1 = getMemoryViewForm().DisassemblerView.SelectedAddress
	local dv_address2 = getMemoryViewForm().DisassemblerView.SelectedAddress2
	local startAddress = math.min(dv_address1, dv_address2)
	local endAddress = math.max(dv_address1, dv_address2)
	local length = endAddress + getInstructionSize(endAddress) - startAddress
	local position = startAddress
	local signature = ''
	while position < endAddress + getInstructionSize(endAddress) do
		local instLength = getInstructionSize(position)
		local disStr = disassemblerView.disassemble(position)
		local _, _, bytes, _ = splitDisassembledString(disStr)
		local addressWildCardStr = string.rep(wildCard, 8 / #wildCard)
		bytes = bytes:gsub('%x%x%x%x%x%x%x%x%s*$', addressWildCardStr):gsub(' ', '')
		signature = signature..bytes
		position = position + instLength
	end
	return signature:gsub(wildCard..'*$', '')
end

local function getHookNames(templateName, settings)
	templateName = settings.Caption or templateName
	local name = inputQuery(templateName..' Hook Name', 'Enter the hook name', '')
	if not name then return end
	if settings.AppendToHookName then
		if not name:match(settings.AppendToHookName..'$') then
			name = name..' '..settings.AppendToHookName
		end
	end
	local hookName = ''
	local hookNameParsed = ''
	for word in (name..' '):gmatch('(%w*)%s') do
		local fl = word:sub(1, 1):upper()
		local rl = word:sub(2)
		hookName = hookName..fl..rl..' '
		hookNameParsed = hookNameParsed..fl..rl
	end
	hookName = hookName:sub(1, -2)
	return hookName, hookNameParsed
end

local function getModuleName(address)
	local dv_address1 = getMemoryViewForm().DisassemblerView.SelectedAddress
	local dv_address2 = getMemoryViewForm().DisassemblerView.SelectedAddress2
	local startAddress = math.min(dv_address1, dv_address2)
	local endAddress = math.max(dv_address1, dv_address2)
	local moduleName
	if inModule(address) and inModule(startAddress) and inModule(endAddress) then
		moduleName = getNameFromAddress(address)
		moduleName = moduleName:match('^(.-)[+-]%x*$')
		moduleName = getNameFromAddress(address)
	end
	return moduleName
end

local function getInjectionInfo(address, aobAddress, injectionInfoLinesCount)
	address = getAddressSafe(address)
	aobAddress = getAddressSafe(aobAddress)
	injectionInfoLinesCount = injectionInfoLinesCount or 20
	local function getLine(lineAddr)
		local disassembler = getVisibleDisassembler()
		local disassembledStr = disassembler.disassemble(lineAddr):gsub('{.-}', '')
		local extraField, opcode, bytes, address = splitDisassembledString(disassembledStr)
		address = getNameFromAddress(address)
		return address..' - '..
				string.format('%-24s', bytes)..'- '..opcode
	end
	local code = { getLine(address)..'   <<<---- Injection point' }
	local addr = getPreviousOpcode(address)
	for i = 1, injectionInfoLinesCount / 2 do
		local str = getLine(addr)
		if aobAddress and addr == aobAddress then
			str = str..'   <<<---- AOB'
		end
		table.insert(code, 1, str)
		addr = getPreviousOpcode(addr)
	end	
	addr = address + getInstructionSize(address)
	for i = 1, injectionInfoLinesCount / 2 do
		local str = getLine(addr)
		if aobAddress and addr == aobAddress then
			str = str..'   <<<---- AOB'
		end
		table.insert(code, str)
		addr = addr + getInstructionSize(addr)
	end
	return code
end

local function generateScript(name, script, sender)
	local settingsPath = TemplatesFolder..sep..name..'.settings.lua'
	local env = { FinalCompilation = false }
	env.__index = env
	setmetatable(env, { __index = _G })
	local settings = exists(settingsPath) and loadfile(settingsPath, 't', env)() or { }

	local dv_address1 = getMemoryViewForm().DisassemblerView.SelectedAddress
	local dv_address2 = getMemoryViewForm().DisassemblerView.SelectedAddress2
	local startAddress = math.min(dv_address1, dv_address2)
	local endAddress = math.max(dv_address1, dv_address2)
	local addressString = getNameFromAddress(startAddress)
	local address = startAddress
	if settings.AskForInjectionAddress then
		address = inputQuery((settings.Caption or name)..' Injection Point', 
								'Enter the injection point address', addressString)
		if not address then return end
		address = getAddress(address)
	end
	local moduleName = getModuleName(address)
	local hookName
	local hookNameParsed
	if settings.AskForHookName then
		hookName, hookNameParsed = getHookNames(name, settings)
		if not hookName then return end
	end
	local aobOffset = startAddress - address
	local aobOffsetString
	if aobOffset < 0 then
		aobOffsetString = '+'..string.format('%X', math.abs(aobOffset))
	elseif aobOffset > 0 then
		aobOffsetString = '-'..string.format('%X', math.abs(aobOffset))
	end

	local position = address
	local injectionLength = 0
	local originalBytes = { }
	local originalBytesString = ''
	local baseAddressRegistry
	local baseAddressOffset
	local originalOpcode = { }
	local vdv = getVisibleDisassembler()
	local ddv = getDefaultDisassembler()
	while injectionLength < 5 do
		local instLength = getInstructionSize(position)
		local disStr = vdv.disassemble(position)
		local _, opcode, _, _ = splitDisassembledString(disStr)
		if not baseAddressRegistry then
			baseAddressRegistry = opcode:match('%[{R}(%w-){N}') or ''
			baseAddressOffset = (opcode:match('{N}([%-%+]{H}%x-){N}%]') or ''):gsub('{H}', '')
		end
		local disStr = ddv.disassemble(position)
		local _, opcode, _, _ = splitDisassembledString(disStr)
		table.insert(originalOpcode, opcode)
		local bytes = readBytes(position, instLength, true)
		for _, b in ipairs(bytes or {}) do
			table.insert(originalBytes, b)
			originalBytesString = originalBytesString..string.format('%02X ', b)
		end
		injectionLength = injectionLength + instLength
		position = position + instLength
	end
	originalBytesString = originalBytesString:sub(1, -2)

	local nopBytes = { }
	local nopBytesString = ''
	if injectionLength > 5 then
		for i = 6, injectionLength do
			table.insert(nopBytes, 0x90)
			nopBytesString = nopBytesString..'90 '
		end
		nopBytesString = nopBytesString:sub(1, -2)
	end

	local env = {
		HookName = hookName, 
		HookNameParsed = hookNameParsed, 
		ModuleName = moduleName, 
		Address = address, 
		AddressString = getNameFromAddress(address), 
		AobAddress = startAddress, 
		AobAddressString = getNameFromAddress(startAddress), 
		OriginalBytes = originalBytes, 
		OriginalBytesString = originalBytesString, 
		Signature = generateSignature(settings.WildCard), 
		AOB = readBytes(startAddress, endAddress - startAddress, true), 
		AobOffset = aobOffset, 
		AobOffsetString = aobOffsetString, 
		PointerDefault = targetIs64Bit() and 'dq 0' or 'dd 0', 
		BaseAddressRegistry = baseAddressRegistry, 
		BaseAddressOffset = baseAddressOffset, 
		OriginalOpcode = originalOpcode, 
		InjectionInfo = getInjectionInfo(address, startAddress, settings.InjectionInfoLinesCount), 
		NopBytes = nopBytes, 
		NopBytesString = nopBytesString, 
		CEVersion = getCEVersion(), 
		Date = os.date('%x'), 
		FinalCompilation = true, 
	}
	env.__index = env
	setmetatable(env, { __index = _G })
	if exists(settingsPath) then
		settings = loadfile(settingsPath, 't', env)() or { }
	end
	local path = TemplatesFolder..sep..'Header.'..CeaFileExtension
	if exists(path) then
		local template, err = TemplateEngine.compileFile(path, env)
		if template and not err then
			env.Header = template
		else
			print(err)
		end
	end
	path = TemplatesFolder..sep..name..'.'..CeaFileExtension
	template, err = TemplateEngine.compileFile(path, env)
	if template and not err then
		if type(script.addText) == 'function' then
			script.addText(template)
		else
			local s = script.getText()
			script.clear()
			script.setText(s..template)
		end
		if type(createMemoryRecords) == 'function' and settings.MemoryRecords then
			env.Script = template
			setmetatable(env, { })
			createMemoryRecords(settings.MemoryRecords, env, false)
		end
	else
		print(err)
	end
end

local function loadTemplates(thread)
	local fileList
	if exists(TemplatesFolder..sep..'LoadOrder.lua') then
		fileList = dofile(TemplatesFolder..sep..'LoadOrder.lua')
	end
	if not fileList then fileList = getFileList(TemplatesFolder) end
	for i, fp in ipairs(fileList) do
		local fn = fp:match('([^'..sep..']+)$')
		local fn = fn:match('^(.+)%..-$')
		if fn ~= 'Header' and fn ~= 'LoadOrder' and not fn:match('.settings$') then
			local settingsPath = TemplatesFolder..sep..fn..'.settings.lua'
			local env = { FinalCompilation = false }
			env.__index = env
			setmetatable(env, { __index = _G })
			local settings = exists(settingsPath) and loadfile(settingsPath, 't', env)() or { }
			registerAutoAssemblerTemplate(settings.Caption or fn, function(script, sender)
				generateScript(fn, script, sender)
			end, settings.Shortcut)
		end
	end
	thread.terminate()
end

createThread(loadTemplates)