import gradio as gr

from llama_cpp import Llama

llm = Llama.from_pretrained(
    repo_id="ID2223JR/gguf_model",
    filename="unsloth.Q4_K_M.gguf",
)


# llm = Llama.from_pretrained(
#    repo_id="ID2223JR/gguf_model_q8",
#    filename="unsloth.Q8_0.gguf",
# )


# Data storage
ingredients_list = []


# Function to add ingredient
def add_ingredient(ingredient, quantity):
    if ingredient and quantity > 0:
        ingredients_list.append(f"{ingredient}, {quantity} grams")
    return (
        "\n".join(ingredients_list),
        gr.update(value="", interactive=True),
        gr.update(value=None, interactive=True),
    )


# Function to enable/disable add button
def validate_inputs(ingredient, quantity):
    if ingredient and quantity is not None and quantity > 0:
        return gr.update(interactive=True)
    return gr.update(interactive=False)


def submit_to_model():
    if not ingredients_list:
        yield "Ingredients list is empty! Please add ingredients first."
        return

    prompt = f"Using the following ingredients, suggest a recipe:\n\n" + "\n".join(
        ingredients_list
    )

    try:
        response = llm.create_chat_completion(
            messages=[
                {
                    "role": "system",
                    "content": (
                        "You are a world-renowned chef, celebrated for your expertise..."
                    ),
                },
                {"role": "user", "content": prompt},
            ],
            stream=True,  # Enable streaming
        )

        content = ""

        for partial_response in response:

            content += partial_response["choices"][0]["delta"].get("content", "")

            if content:
                yield content

        ingredients_list.clear()  # Reset list after generation

    except Exception as e:
        yield f"An error occurred: {str(e)}"


# App
def app():
    with gr.Blocks() as demo:
        with gr.Row():
            ingredient_input = gr.Textbox(
                label="Ingredient", placeholder="Enter ingredient name"
            )
            quantity_input = gr.Number(label="Quantity (grams)", value=None)

        add_button = gr.Button("Add Ingredient", interactive=False)
        output = gr.Textbox(label="Ingredients List", lines=10, interactive=False)

        with gr.Row():
            submit_button = gr.Button("Submit")
            model_output = gr.Textbox(
                label="Recipe Suggestion", lines=25, interactive=False
            )

        # Validate inputs
        ingredient_input.change(
            validate_inputs, [ingredient_input, quantity_input], add_button
        )
        quantity_input.change(
            validate_inputs, [ingredient_input, quantity_input], add_button
        )

        # Add ingredient logic
        add_button.click(
            add_ingredient,
            [ingredient_input, quantity_input],
            [output, ingredient_input, quantity_input],
        )

        # Submit to model logic
        submit_button.click(
            submit_to_model,
            inputs=None,  # No inputs required as it uses the global ingredients_list
            outputs=model_output,
        )

    return demo


demo = app()
demo.launch()