from threading import Lock

import execution

CPACK_LAST_ID = None
CPACK_OUTPUT_CACHE = {}

_lock = Lock()


def store_bentoml_value(func):
    def wrapped(
        inputs, class_def, unique_id, outputs=None, dynprompt=None, extra_data={}
    ):
        global CPACK_LAST_ID
        if getattr(class_def, "CPACK_NODE", False):
            with _lock:
                CPACK_LAST_ID = unique_id
        if outputs is None:
            outputs = CPACK_OUTPUT_CACHE
        return func(inputs, class_def, unique_id, outputs, dynprompt, extra_data)

    return wrapped


execution.get_input_data = store_bentoml_value(execution.get_input_data)


def set_bentoml_output(output):
    with _lock:
        CPACK_OUTPUT_CACHE[CPACK_LAST_ID] = output