import gradio as gr
import requests
import os
from PIL import Image
import io
import base64
import json
import tempfile

class FashionClassifier:
    def __init__(self):
        self.api_url = "https://api.marqo.ai/classify"
        self.api_key = os.getenv("MARQO_API_KEY")
        
        # Vérifier si la clé API est configurée
        if not self.api_key or self.api_key == "your_marqo_api_key":
            print("⚠️ ATTENTION: Clé API Marqo non configurée!")
            print("👉 Ajoutez MARQO_API_KEY dans les secrets Hugging Face")

    def classify_image(self, image, max_categories=5, confidence_threshold=0.3):
        """Classifie une image avec gestion d'erreurs complète"""
        try:
            # Vérifier si l'image est valide
            if image is None:
                return {"error": "Aucune image fournie"}
            
            # Convertir et optimiser l'image
            image = self.prepare_image(image)
            
            # Convertir en base64
            buffered = io.BytesIO()
            image.save(buffered, format="JPEG", optimize=True, quality=85)
            img_str = base64.b64encode(buffered.getvalue()).decode()
            
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.api_key}"
            }
            
            payload = {
                "model": "Marqo/Marqo-FashionSigLIP-Classification",
                "image_data": img_str,
                "parameters": {
                    "max_categories": int(max_categories),
                    "confidence_threshold": float(confidence_threshold)
                }
            }
            
            # Debug: Afficher la requête (dans les logs)
            print(f"📤 Envoi requête à Marqo API...")
            
            response = requests.post(self.api_url, headers=headers, json=payload, timeout=30)
            
            # Debug: Afficher la réponse
            print(f"📥 Réponse reçue: {response.status_code}")
            
            if response.status_code != 200:
                return {
                    "error": f"Erreur API {response.status_code}",
                    "details": response.text
                }
                
            return response.json()
            
        except Exception as e:
            return {"error": f"Erreur: {str(e)}"}

    def prepare_image(self, image):
        """Prépare l'image pour le modèle"""
        # Convertir en RGB si nécessaire
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        # Redimensionner si trop grande (max 1024px)
        max_size = (1024, 1024)
        image.thumbnail(max_size, Image.Resampling.LANCZOS)
        
        return image

# Initialiser le classifieur
classifier = FashionClassifier()

def process_image(image, max_categories, confidence):
    """Fonction de traitement avec debug"""
    try:
        if image is None:
            return "❌ Veuillez uploader une image valide"
        
        print(f"🖼️ Image reçue: {image.size}, mode: {image.mode}")
        
        result = classifier.classify_image(image, max_categories, confidence)
        
        # Debug dans les logs
        print(f"🔍 Résultat brut: {json.dumps(result, indent=2)}")
        
        if "error" in result:
            error_msg = f"❌ Erreur de classification:\n\n"
            error_msg += f"**Détails:** {result['error']}\n\n"
            
            # Conseils selon le type d'erreur
            if "401" in str(result['error']):
                error_msg += "🔑 **Solution:** Vérifiez votre clé API Marqo dans les secrets Hugging Face"
            elif "timeout" in str(result['error']).lower():
                error_msg += "⏱️ **Solution:** L'image est trop lourde, essayez une image plus petite"
            else:
                error_msg += "💡 **Solution:** Essayez une image plus petite ou un format différent"
                
            return error_msg
        
        if "predictions" in result and result["predictions"]:
            output = "## 🎯 Résultats de classification:\n\n"
            for i, pred in enumerate(result["predictions"]):
                output += f"{i+1}. **{pred['label']}** - {pred['score']*100:.1f}%\n"
            
            if "processing_time" in result:
                output += f"\n⏱️ Temps de traitement: {result['processing_time']}s"
            
            return output
        else:
            return "❌ Aucune prédiction trouvé - Le modèle n'a pas reconnu l'image"
            
    except Exception as e:
        return f"❌ Erreur inattendue: {str(e)}"

# Interface Gradio simplifiée et robuste
with gr.Blocks(title="Classificateur de Mode") as demo:
    gr.Markdown("""
    # 🎨 Classificateur de Vêtements Marqo
    **Uploader une image de vêtement** pour la classifier automatiquement
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Uploader votre image")
            image_input = gr.Image(
                type="pil", 
                label="Image à classifier",
                height=300,
                sources=["upload", "webcam", "clipboard"]
            )
            
            gr.Markdown("### ⚙️ Paramètres")
            max_categories = gr.Slider(1, 10, value=5, label="Catégories max")
            confidence = gr.Slider(0.1, 1.0, value=0.3, label="Seuil de confiance")
            
            gr.Markdown("### 🚀 Actions")
            submit_btn = gr.Button("Classifier l'image", variant="primary")
            clear_btn = gr.Button("Effacer", variant="secondary")
        
        with gr.Column(scale=2):
            gr.Markdown("### 📊 Résultats")
            output_text = gr.Markdown(
                value="⬅️ Uploader une image de vêtement pour commencer"
            )
    
    # Exemples d'images qui FONCTIONNENT
    gr.Examples(
        examples=[
            ["https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png"],
            ["https://images.unsplash.com/photo-1558769132-cb1aea458c5e"],
            ["https://images.unsplash.com/photo-1543163521-1bf539c55dd2"]
        ],
        inputs=image_input,
        label="🖼️ Exemples testés et fonctionnels"
    )
    
    # Événements
    submit_btn.click(
        fn=process_image,
        inputs=[image_input, max_categories, confidence],
        outputs=output_text
    )
    
    clear_btn.click(
        fn=lambda: (None, 5, 0.3, "⬅️ Uploader une nouvelle image"),
        inputs=[],
        outputs=[image_input, max_categories, confidence, output_text]
    )

# Lancer avec debug
if __name__ == "__main__":
    demo.launch(debug=True)