from beakers.recipe import Recipe
import pydantic


class Word(pydantic.BaseModel):
    word: str


class ClassifiedWord(pydantic.BaseModel):
    normalized_word: str
    is_fruit: bool


class Sentence(pydantic.BaseModel):
    sentence: list[str]


def word_classifier(item) -> ClassifiedWord:
    return ClassifiedWord(
        normalized_word=item.word.lower(),
        is_fruit=item.word.lower()
        in (
            "apple",
            "banana",
            "fig",
            "grape",
            "lemon",
            "mango",
            "orange",
            "pear",
            "raspberry",
        ),
    )


recipe = Recipe("fruits-example")
recipe.add_beaker("word", Word)
recipe.add_beaker("classified_word", ClassifiedWord)
recipe.add_beaker("sentence", Sentence)
recipe.add_transform("word", "classified_word", word_classifier)
recipe.add_conditional(
    "classified_word",
    lambda cw: cw.is_fruit,
    "fruits",
)
recipe.add_transform("fruits", "sentence", lambda x: f"I love a fresh {x['word']}")

recipe.add_seed(
    "word",
    [
        Word(word="apple"),
        Word(word="banana"),
        Word(word="hammer"),
        Word(word="orange"),
        Word(word="egg"),
    ],
)