from transformers import AutoProcessor, AutoModelForImageClassification, pipeline
from PIL import Image
import torch


def load_classification_model():
    processor = AutoProcessor.from_pretrained("Shresthadev403/food-image-classification")
    model = AutoModelForImageClassification.from_pretrained("Shresthadev403/food-image-classification")
    return processor, model
def load_text_generator():
    return pipeline( "text2text-generation",model="distilgpt2")
processor, model = load_classification_model()
text_generator = load_text_generator()


def predict_dish(image: Image.Image):
    try:
        image = Image.open(image).convert("RGB")
        inputs = processor(images=image, return_tensors="pt")
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            predicted_class_idx = logits.argmax(-1).item()
            label = model.config.id2label[predicted_class_idx]
            return label.lower().replace(" ", "_")
    except Exception as e:
        print(f"❌ Error in dish prediction: {e}")
        return "unknown_dish"



def generate_recipe(dish, diet=None, cuisine=None, cook_time=None):

    filters = []
    if diet and diet != "Any":
        filters.append(f"{diet} diet")
    if cuisine and cuisine != "Any":
        filters.append(f"{cuisine} cuisine")
    if cook_time and cook_time != "Any":
        filters.append(f"ready in {cook_time}")
        filter_text = ", ".join(filters)
    
    prompt = f"""
    Create a step-by-step recipe for {dish}.
    Include:
    - Ingredients with quantities
    - Step-by-step instructions cooking steps
    Make sure it's a {filter_text} recipe."""
    try:
        result = text_generator(prompt.strip(), max_length=282, do_sample=False)
        return result[0]['generated_text']
    except Exception as e:
        print(f"❌ Error generating recipe: {e}")
        return "Sorry, couldn't generate a recipe at the moment."