from typing import Literal
import pydantic
from prefect import flow, pause_flow_run, runtime
from prefect.input import RunInput


class DinnerMenu(RunInput):
    eat_out: Literal["yes", "YES!!"]
    menu: Literal["Burger", "Sandwich", "Vietnamese", "Indian", "Cook in home"]

    @pydantic.validator("menu")
    def validate_age(cls, value, values, **kwargs):
        if value == "Cook in home" :
            raise ValueError("We need to eat out today. Too tired to cook sth.")
        return value

def success_hook(flow, flow_run, state):
    print(f"Flow : {flow_run.id} succeeded!")

@flow(on_completion=[success_hook])
def get_dinner_menu():
    dinner_menu = pause_flow_run(wait_for_input=DinnerMenu)
    print(f"Flow {runtime.flow_run.name} picked dinner menu {dinner_menu} !")


if __name__ == "__main__":
    get_dinner_menu()