from java.util import ArrayList

from ghidra.program.model.data import DWordDataType, WordDataType
from ghidra.program.model.listing import Function, ParameterImpl, ReturnParameterImpl, VariableStorage
from ghidra.program.model.symbol import SourceType


def _clone(data_type):
    return data_type.clone(program.getDataTypeManager())


function = helpers["get_function"](program, "1000:42e2")
if function is None:
    raise RuntimeError("Function 1000:42e2 not found")

dword_type = _clone(DWordDataType.dataType)
word_type = _clone(WordDataType.dataType)

param_1 = ParameterImpl("base_ptr", dword_type, VariableStorage(program, 4, 4), program, SourceType.USER_DEFINED)
blocksize = ParameterImpl("blocksize", word_type, VariableStorage(program, 8, 2), program, SourceType.USER_DEFINED)
nblocks = ParameterImpl("nblocks", word_type, VariableStorage(program, 10, 2), program, SourceType.USER_DEFINED)
param_4 = ParameterImpl("param_4", word_type, VariableStorage(program, 12, 2), program, SourceType.USER_DEFINED)
param_5 = ParameterImpl("param_5", word_type, VariableStorage(program, 14, 2), program, SourceType.USER_DEFINED)
transform_func = ParameterImpl("transform_func", dword_type, VariableStorage(program, 16, 4), program, SourceType.USER_DEFINED)

ax_reg = program.getRegister("AX")
dx_reg = program.getRegister("DX")
return_param = ReturnParameterImpl(dword_type, VariableStorage(program, ax_reg, dx_reg), program)

params = ArrayList()
params.add(param_1)
params.add(blocksize)
params.add(nblocks)
params.add(param_4)
params.add(param_5)
params.add(transform_func)

function.updateFunction(
    function.getCallingConventionName(),
    return_param,
    params,
    Function.FunctionUpdateType.CUSTOM_STORAGE,
    True,
    SourceType.USER_DEFINED,
)

print("updated", function.getEntryPoint(), function.getSignature())
for parameter in function.getParameters():
    print("param", parameter.getName(), parameter.getDataType().getDisplayName(), parameter.getVariableStorage())
print("return", function.getReturnType().getDisplayName(), function.getReturn().getVariableStorage())