import gradio as gr

from transformers import TextStreamer

from unsloth import FastLanguageModel

max_seq_length = 2048
dtype = None
load_in_4bit = True

peft_model_id = "ID2223JR/lora_model"

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=peft_model_id,
    max_seq_length=max_seq_length,
    dtype=dtype,
    load_in_4bit=load_in_4bit,
)
FastLanguageModel.for_inference(model)


# Data storage
ingredients_list = []


# Function to add ingredient
def add_ingredient(ingredient, quantity):
    if ingredient and int(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 > 0:
        return gr.update(interactive=True)
    return gr.update(interactive=False)


# Function to handle model submission
def submit_to_model():
    if not ingredients_list:
        return "Ingredients list is empty! Please add ingredients first."

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

    messages = [
        {
            "role": "system",
            "content": "You are a world-renowned chef, celebrated for your expertise in creating delectable dishes from diverse cuisines. You have a vast knowledge of ingredients, cooking techniques, and dietary preferences. Your role is to suggest personalized recipes based on the ingredients available, dietary restrictions, or specific meal requests. Please provide clear, step-by-step instructions and any useful tips to enhance the dish's flavor or presentation. Begin by introducing the recipe and why it’s a great choice.",
        },
        {"role": "user", "content": prompt},
    ]
    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,  # Must add for generation
        return_tensors="pt",
    )
    text_streamer = TextStreamer(tokenizer, skip_prompt=True)

    return model.generate(
        input_ids=inputs,
        streamer=text_streamer,
        max_new_tokens=128,
        use_cache=True,
        temperature=1.5,
        min_p=0.1,
    )


# 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=10, 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()