import json
import sys

sys.path.append("/DIRTY/dirty")

import utils.infer
from utils.dataset import Example


def dump_functions_to_json():
    functionManager = currentProgram().getFunctionManager()
    functions = functionManager.getFunctions(True)
    
    # Create a dictionary to store function names and their corresponding addresses
    functions_dict = {}
    
    for func in functions:

        if func.isExternal() or func.isThunk():
            continue

        func_name = func.getName()
        func_address = func.getEntryPoint().getOffset()
        
        cf = utils.infer.ghidra_obtain_cf(func)

        example = Example.from_cf(
            cf, binary_file="binary_file", max_stack_length=1024, max_type_size=1024
        )

        # Add function name and address to the dictionary
        if example.is_valid_example:
            functions_dict[func_address] = (func_name, cf.to_json(), len(example.source))
    
    # Convert the dictionary to a JSON object
    json_data = json.dumps(functions_dict, indent=4)
    
    args = getScriptArgs()

    if len(args) == 0:

        # Print JSON output to the console
        print(json_data)

    else:
            
        # Write JSON output to a file
        with open(args[0], 'w') as f:
            f.write(json_data)
    
# Run the function
dump_functions_to_json()