# ===========================
# Recipe Chatbot for Hugging Face Space
# ===========================

import gradio as gr
import os

# ---------------------------
# Install transformers if not installed (optional in Spaces)
# ---------------------------
try:
    from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
except ModuleNotFoundError:
    import subprocess
    subprocess.check_call(["pip", "install", "transformers", "torch", "gradio"])
    from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# ---------------------------
# Define model folder
# ---------------------------
# Make sure your folder in the Space is exactly named 'recipe-model'
MODEL_PATH = "recipe-model"  # Do NOT use './recipe-model'

if not os.path.exists(MODEL_PATH):
    raise FileNotFoundError(f"Model folder '{MODEL_PATH}' not found. Upload your trained GPT-2 model.")

# ---------------------------
# Load tokenizer and model
# ---------------------------
try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    model = AutoModelForCausalLM.from_pretrained(MODEL_PATH)
    recipe_generator = pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        device=-1  # use CPU; set 0 if GPU available
    )
except Exception as e:
    raise RuntimeError(f"Failed to load model: {e}")

# ---------------------------
# Chatbot function
# ---------------------------
def get_recipe(user_input):
    """
    Takes ingredients as input and returns a generated recipe.
    """
    if not user_input.strip():
        return "Please enter some ingredients."

    prompt = f"Recipes with {user_input}:"
    try:
        result = recipe_generator(
            prompt,
            max_length=250,  # Adjust for longer recipes
            num_return_sequences=1,
            do_sample=True,
            temperature=0.7,
            top_p=0.9
        )
        # Remove the prompt from the output for clean response
        generated_text = result[0]["generated_text"]
        if generated_text.lower().startswith(prompt.lower()):
            generated_text = generated_text[len(prompt):].strip()
        return generated_text
    except Exception as e:
        return f"Error generating recipe: {e}"

# ---------------------------
# Build Gradio Interface
# ---------------------------
iface = gr.Interface(
    fn=get_recipe,
    inputs=gr.Textbox(
        lines=2,
        placeholder="Enter ingredients (e.g., potato, chicken, cheese)",
        label="Ingredients"
    ),
    outputs=gr.Textbox(
        label="Generated Recipe"
    ),
    title="Recipe Chatbot",
    description="Enter ingredients you have and get a recipe generated by your trained GPT-2 model.",
    examples=[
        ["potato, cheese"],
        ["chicken, rice, onion"],
        ["tomato, basil, mozzarella"]
    ],
    theme="default"
)

# Launch the app (Hugging Face Spaces will run this automatically)
iface.launch()