import inspect
import functools
from prefect import flow, task

USE_PREFECT = True


def custom_task(function):
    def wrapper(*args, **kwargs):
        print("Task wrapper started")
        ret = function(*args, **kwargs)
        print("Task wrapper ended")
        return ret

    return wrapper


task_decorator = task if USE_PREFECT else custom_task


def do_something(function):
    signature = inspect.signature(function)

    @functools.wraps(function)
    def wrapper(*args, **kwargs):
        print(args, kwargs)

        return function(**kwargs)

    wrapper.__signature__ = signature

    return wrapper


@task_decorator
@do_something
def normal_task(message: str, status_code: int):
    print(f"{message=}")
    print(f"{status_code=}")


@flow
def entry_flow():
    normal_task(message="Hello", status_code=200)


if __name__ == "__main__":
    entry_flow()