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, "1420:1499")
if function is None:
    raise RuntimeError("Function 1420:1499 not found")

dword_type = _clone(DWordDataType.dataType)
word_type = _clone(WordDataType.dataType)
runtime_offset_param = ParameterImpl("this", word_type, VariableStorage(program, 4, 2), program, SourceType.USER_DEFINED)
runtime_segment_param = ParameterImpl("runtime_segment", word_type, VariableStorage(program, 6, 2), program, SourceType.USER_DEFINED)
owner_type_param = ParameterImpl("owner_type", word_type, VariableStorage(program, 8, 2), program, SourceType.USER_DEFINED)
owner_id_param = ParameterImpl("owner_id", word_type, VariableStorage(program, 10, 2), 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(runtime_offset_param)
params.add(runtime_segment_param)
params.add(owner_type_param)
params.add(owner_id_param)

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

function.setName("Create", SourceType.USER_DEFINED)
function.setComment("Factory-style runtime creator. Uses split 16-bit this/segment parameters so Ghidra can represent the incoming far runtime pointer without corrupting decompilation.")

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