import gradio as gr
from transformers import pipeline
from collections import Counter
import re

# 1. Load an instruction-tuned model for text generation
generator = pipeline(
    "text2text-generation",
    model="google/flan-t5-small",
    device=0  # set to -1 for CPU
)

def plan_and_list(ingredients, days, goal):
    # 2. Build the prompt
    prompt = (
        f"Plan {days} days of breakfast, lunch, and dinner menus "
        f"using these ingredients: {ingredients}. "
        f"Meet {goal}."
    )
    # 3. Generate the meal plan
    result = generator(prompt, max_length=512, do_sample=False)[0]["generated_text"]
    
    # 4. Parse meals and build grocery list
    #    (naïve parsing—split lines, extract items after colon)
    grocery_counter = Counter()
    for line in result.split("\n"):
        # e.g. "Day 1 Lunch: Grilled chicken with quinoa and spinach"
        if ":" in line and any(meal in line for meal in ["Breakfast", "Lunch", "Dinner"]):
            items = re.split(r"with|and|,", line.split(":",1)[1])
            for item in items:
                item = item.strip().lower()
                if item:
                    grocery_counter[item] += 1
    
    shopping_list = "\n".join(f"{item}: {qty}" for item, qty in grocery_counter.items())
    return result, shopping_list

# 5. Gradio UI
iface = gr.Interface(
    fn=plan_and_list,
    inputs=[
        gr.Textbox(label="Ingredients (comma-separated)", placeholder="e.g. eggs, spinach, quinoa…"),
        gr.Slider(minimum=1, maximum=7, step=1, label="Days to plan"),
        gr.Textbox(label="Calorie Goal / Notes", placeholder="e.g. 1800 kcal/day, vegetarian…")
    ],
    outputs=[
        gr.Textbox(label="Generated Meal Plan"),
        gr.Textbox(label="Consolidated Grocery List")
    ],
    title="🍽️ Smart Meal Planner & Grocery List Generator",
    description="Enter what you have (or your dietary goals) and get a weekly menu plus a unified shopping list.",
)
iface.launch()