import os
import requests
import gradio as gr
from transformers import ViTFeatureExtractor, ViTForImageClassification, T5Tokenizer, T5ForConditionalGeneration
from PIL import Image
import torch

# Charger le modèle de classification d'image (ViT)
vit_model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')

# Charger le modèle de génération de texte (T5)
t5_model = T5ForConditionalGeneration.from_pretrained('flax-community/t5-recipe-generation')
t5_tokenizer = T5Tokenizer.from_pretrained('flax-community/t5-recipe-generation')

# Récupérer la clé API depuis les secrets
api_key = os.getenv('GOOGLE_API_KEY')

# Exemple de fonction pour utiliser l'API avec la clé
def call_external_api(query):
    headers = {"Authorization": f"Bearer {api_key}"}
    response = requests.get(f"https://api.example.com/search?query={query}", headers=headers)
    return response.json()

# Fonction pour prédire l'aliment à partir d'une image
def classify_image(image):
    inputs = feature_extractor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = vit_model(**inputs)
    predicted_class_idx = outputs.logits.argmax(-1).item()
    predicted_class = vit_model.config.id2label.get(predicted_class_idx, "Inconnu")
    
    return predicted_class

# Fonction pour générer une recette à partir du nom de l'aliment
def generate_recipe(food_name):
    # Exemple d'utilisation de l'API avec la clé
    api_response = call_external_api(food_name)
    # Supposons que l'API retourne directement la recette ou que tu la traites ici
    recipe = api_response.get('recipe', 'Recette non trouvée')
    
    return recipe

# Pipeline complet : image -> nom de l'aliment -> recette
def predict_and_generate(image):
    # Prédire l'aliment
    food_name = classify_image(image)
    
    if food_name == "Inconnu":
        return "Aucun aliment reconnu", "Essayez avec une autre image."
    
    # Générer la recette
    recipe = generate_recipe(food_name)
    
    return food_name, recipe

# Interface Gradio améliorée
interface = gr.Interface(
    fn=predict_and_generate,
    inputs=gr.Image(type="pil"),
    outputs=["text", "text"],
    title="Générateur de recettes à partir d'images d'aliments",
    description="Téléchargez une image d'un aliment, cet outil l'identifiera et générera une recette correspondante."
)

interface.launch()