import os
import gradio as gr
import base64
import io
import json
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from typing import List, Dict, Any, Tuple
import requests
from sentence_transformers import SentenceTransformer
import faiss
from groq import Groq
import tempfile
import re

class FoodRAGApplication:
    def __init__(self):
        """Initialize the Food Recognition RAG application"""
        # Initialize Groq client for vision and text
        self.groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
        
        # Initialize embedding model for nutrition database
        self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
        
        # Initialize FAISS index for nutrition knowledge
        self.dimension = 384
        self.nutrition_index = faiss.IndexFlatIP(self.dimension)
        
        # Comprehensive nutrition database
        self.nutrition_database = self._create_nutrition_database()
        self.nutrition_embeddings = []
        self.is_nutrition_indexed = False
        
        # Build nutrition index
        self._build_nutrition_index()
    
    def _create_nutrition_database(self) -> List[Dict]:
        """Create a comprehensive nutrition database for common foods"""
        nutrition_db = [
            # Fruits
            {
                "food_name": "Apple",
                "category": "Fruit",
                "calories_per_100g": 52,
                "carbs": 14,
                "fiber": 2.4,
                "sugar": 10,
                "protein": 0.3,
                "fat": 0.2,
                "vitamin_c": 4.6,
                "potassium": 107,
                "origin": "Central Asia",
                "season": "Fall",
                "health_benefits": "Rich in antioxidants, supports heart health, aids digestion",
                "description": "Crisp, sweet fruit high in fiber and vitamin C"
            },
            {
                "food_name": "Banana",
                "category": "Fruit",
                "calories_per_100g": 89,
                "carbs": 23,
                "fiber": 2.6,
                "sugar": 12,
                "protein": 1.1,
                "fat": 0.3,
                "vitamin_c": 8.7,
                "potassium": 358,
                "origin": "Southeast Asia",
                "season": "Year-round",
                "health_benefits": "High in potassium, supports muscle function, quick energy source",
                "description": "Tropical fruit rich in potassium and natural sugars"
            },
            {
                "food_name": "Orange",
                "category": "Fruit",
                "calories_per_100g": 47,
                "carbs": 12,
                "fiber": 2.4,
                "sugar": 9,
                "protein": 0.9,
                "fat": 0.1,
                "vitamin_c": 53.2,
                "potassium": 181,
                "origin": "China",
                "season": "Winter",
                "health_benefits": "Excellent source of vitamin C, boosts immunity, supports skin health",
                "description": "Citrus fruit packed with vitamin C and folate"
            },
            {
                "food_name": "Strawberry",
                "category": "Fruit",
                "calories_per_100g": 32,
                "carbs": 8,
                "fiber": 2,
                "sugar": 4.9,
                "protein": 0.7,
                "fat": 0.3,
                "vitamin_c": 58.8,
                "potassium": 153,
                "origin": "Europe and North America",
                "season": "Spring-Summer",
                "health_benefits": "High in antioxidants, supports brain health, anti-inflammatory",
                "description": "Sweet berry rich in vitamin C and antioxidants"
            },
            {
                "food_name": "Grapes",
                "category": "Fruit",
                "calories_per_100g": 62,
                "carbs": 16,
                "fiber": 0.9,
                "sugar": 16,
                "protein": 0.6,
                "fat": 0.2,
                "vitamin_c": 3.2,
                "potassium": 191,
                "origin": "Middle East",
                "season": "Late summer-Fall",
                "health_benefits": "Contains resveratrol, supports heart health, antioxidant properties",
                "description": "Sweet fruit rich in natural sugars and antioxidants"
            },
            
            # Vegetables
            {
                "food_name": "Carrot",
                "category": "Vegetable",
                "calories_per_100g": 41,
                "carbs": 10,
                "fiber": 2.8,
                "sugar": 4.7,
                "protein": 0.9,
                "fat": 0.2,
                "vitamin_c": 5.9,
                "potassium": 320,
                "vitamin_a": 835,
                "origin": "Afghanistan",
                "season": "Fall-Winter",
                "health_benefits": "High in beta-carotene, supports eye health, immune function",
                "description": "Root vegetable rich in beta-carotene and fiber"
            },
            {
                "food_name": "Broccoli",
                "category": "Vegetable",
                "calories_per_100g": 34,
                "carbs": 7,
                "fiber": 2.6,
                "sugar": 1.5,
                "protein": 2.8,
                "fat": 0.4,
                "vitamin_c": 89.2,
                "potassium": 316,
                "origin": "Mediterranean",
                "season": "Fall-Spring",
                "health_benefits": "High in vitamin C and K, supports bone health, cancer-fighting compounds",
                "description": "Cruciferous vegetable packed with vitamins and minerals"
            },
            {
                "food_name": "Spinach",
                "category": "Vegetable",
                "calories_per_100g": 23,
                "carbs": 3.6,
                "fiber": 2.2,
                "sugar": 0.4,
                "protein": 2.9,
                "fat": 0.4,
                "vitamin_c": 28.1,
                "potassium": 558,
                "iron": 2.7,
                "origin": "Persia",
                "season": "Spring-Fall",
                "health_benefits": "High in iron and folate, supports blood health, rich in antioxidants",
                "description": "Leafy green vegetable high in iron and vitamins"
            },
            {
                "food_name": "Tomato",
                "category": "Vegetable",
                "calories_per_100g": 18,
                "carbs": 3.9,
                "fiber": 1.2,
                "sugar": 2.6,
                "protein": 0.9,
                "fat": 0.2,
                "vitamin_c": 13.7,
                "potassium": 237,
                "origin": "South America",
                "season": "Summer",
                "health_benefits": "Rich in lycopene, supports heart health, anti-cancer properties",
                "description": "Versatile fruit-vegetable rich in lycopene and vitamin C"
            },
            {
                "food_name": "Bell Pepper",
                "category": "Vegetable",
                "calories_per_100g": 31,
                "carbs": 7,
                "fiber": 2.5,
                "sugar": 4.2,
                "protein": 1,
                "fat": 0.3,
                "vitamin_c": 127.7,
                "potassium": 211,
                "origin": "Central America",
                "season": "Summer-Fall",
                "health_benefits": "Extremely high in vitamin C, supports immune system, antioxidant rich",
                "description": "Colorful vegetable with exceptional vitamin C content"
            },
            
            # Nuts and Seeds
            {
                "food_name": "Almonds",
                "category": "Nut",
                "calories_per_100g": 579,
                "carbs": 22,
                "fiber": 12.5,
                "sugar": 4.4,
                "protein": 21,
                "fat": 50,
                "vitamin_c": 0,
                "potassium": 733,
                "origin": "Middle East",
                "season": "Late summer",
                "health_benefits": "High in healthy fats, supports heart health, good protein source",
                "description": "Tree nut rich in healthy fats, protein, and vitamin E"
            },
            {
                "food_name": "Avocado",
                "category": "Fruit",
                "calories_per_100g": 160,
                "carbs": 9,
                "fiber": 7,
                "sugar": 0.7,
                "protein": 2,
                "fat": 15,
                "vitamin_c": 10,
                "potassium": 485,
                "origin": "South Central Mexico",
                "season": "Year-round",
                "health_benefits": "Rich in monounsaturated fats, supports heart health, nutrient dense",
                "description": "Creamy fruit high in healthy fats and fiber"
            }
        ]
        return nutrition_db
    
    def _build_nutrition_index(self):
        """Build FAISS index for nutrition database"""
        try:
            # Create text descriptions for embedding
            nutrition_texts = []
            for food in self.nutrition_database:
                text = f"{food['food_name']} {food['category']} {food['description']} {food['health_benefits']} {food['origin']}"
                nutrition_texts.append(text)
            
            # Create embeddings
            embeddings = self.embedding_model.encode(nutrition_texts)
            faiss.normalize_L2(embeddings)
            
            # Add to index
            self.nutrition_index.add(embeddings)
            self.nutrition_embeddings = embeddings
            self.is_nutrition_indexed = True
            
        except Exception as e:
            print(f"Error building nutrition index: {e}")
    
    def encode_image_to_base64(self, image: Image.Image) -> str:
        """Convert PIL Image to base64 string"""
        buffered = io.BytesIO()
        image.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        return f"data:image/jpeg;base64,{img_str}"
    
    def identify_food_with_groq(self, image: Image.Image) -> str:
        """Use Groq vision model to identify food in image"""
        try:
            # Convert image to base64
            base64_image = self.encode_image_to_base64(image)
            
            # Call Groq vision API
            completion = self.groq_client.chat.completions.create(
                model="llama-3.2-11b-vision-preview",
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": "Identify the food item(s) in this image. Provide the name of the food, whether it's a fruit, vegetable, or other category. Be specific and concise. If there are multiple food items, list them all."
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": base64_image
                                }
                            }
                        ]
                    }
                ],
                temperature=0.1,
                max_completion_tokens=512,
                top_p=1,
                stream=False,
                stop=None,
            )
            
            return completion.choices[0].message.content
            
        except Exception as e:
            return f"Error identifying food: {str(e)}"
    
    def search_nutrition_info(self, food_identification: str, top_k: int = 3) -> List[Dict]:
        """Search nutrition database for relevant food information"""
        if not self.is_nutrition_indexed:
            return []
        
        try:
            # Create query embedding
            query_embedding = self.embedding_model.encode([food_identification])
            faiss.normalize_L2(query_embedding)
            
            # Search in nutrition index
            scores, indices = self.nutrition_index.search(query_embedding, top_k)
            
            results = []
            for score, idx in zip(scores[0], indices[0]):
                if idx < len(self.nutrition_database):
                    food_info = self.nutrition_database[idx].copy()
                    food_info['similarity_score'] = float(score)
                    results.append(food_info)
            
            return results
            
        except Exception as e:
            print(f"Nutrition search error: {e}")
            return []
    
    def generate_nutrition_response(self, food_identification: str, nutrition_matches: List[Dict]) -> str:
        """Generate comprehensive nutrition response using Groq"""
        try:
            # Prepare nutrition context
            context = ""
            for i, food in enumerate(nutrition_matches):
                context += f"""
Food {i+1}: {food['food_name']} ({food['category']})
- Calories per 100g: {food['calories_per_100g']}
- Carbohydrates: {food['carbs']}g
- Protein: {food['protein']}g
- Fat: {food['fat']}g
- Fiber: {food['fiber']}g
- Vitamin C: {food['vitamin_c']}mg
- Potassium: {food['potassium']}mg
- Origin: {food['origin']}
- Season: {food['season']}
- Health Benefits: {food['health_benefits']}
- Description: {food['description']}
"""
            
            # Create comprehensive prompt
            prompt = f"""Based on the food identification: "{food_identification}" and the following nutrition database information, provide a comprehensive answer about the nutritional content and other relevant information.

Nutrition Database:
{context}

Please provide:
1. Nutritional breakdown (calories, macronutrients, key vitamins/minerals)
2. Health benefits
3. Origin and seasonal information
4. Any interesting facts about the food
5. Serving size recommendations

Make the response informative, engaging, and well-structured. If the identified food matches closely with the database, use that information. If not, provide general nutritional guidance based on the food type identified.
"""

            # Call Groq for response generation
            completion = self.groq_client.chat.completions.create(
                model="llama-3.3-70b-versatile",
                messages=[
                    {
                        "role": "system",
                        "content": "You are a nutrition expert providing detailed, accurate information about foods. Always cite specific nutritional values when available and give practical health advice."
                    },
                    {
                        "role": "user",
                        "content": prompt
                    }
                ],
                temperature=0.3,
                max_completion_tokens=1000,
                top_p=1,
                stream=False,
            )
            
            return completion.choices[0].message.content
            
        except Exception as e:
            return f"Error generating nutrition response: {str(e)}"
    
    def process_food_image(self, image: Image.Image) -> Tuple[str, str, str]:
        """Main function to process food image and return nutrition information"""
        if image is None:
            return "Please upload an image of food.", "", ""
        
        try:
            # Step 1: Identify food using vision model
            food_identification = self.identify_food_with_groq(image)
            
            # Step 2: Search nutrition database
            nutrition_matches = self.search_nutrition_info(food_identification)
            
            # Step 3: Generate comprehensive response
            nutrition_response = self.generate_nutrition_response(food_identification, nutrition_matches)
            
            # Step 4: Create detailed breakdown
            breakdown = "🔍 **Food Identification:**\n"
            breakdown += f"{food_identification}\n\n"
            
            if nutrition_matches:
                breakdown += "📊 **Matching Nutrition Data:**\n"
                for i, food in enumerate(nutrition_matches[:2]):
                    breakdown += f"**{food['food_name']}** ({food['category']})\n"
                    breakdown += f"• Calories: {food['calories_per_100g']} per 100g\n"
                    breakdown += f"• Similarity Score: {food['similarity_score']:.3f}\n\n"
            
            return nutrition_response, breakdown, food_identification
            
        except Exception as e:
            return f"Error processing image: {str(e)}", "", ""

# Initialize the application
food_app = FoodRAGApplication()

def create_food_pattern_background():
    """Create an attractive food pattern background"""
    # Create a large canvas
    width, height = 1200, 800
    background = Image.new('RGB', (width, height), '#f8f9fa')
    draw = ImageDraw.Draw(background)
    
    # Food emojis and colors
    food_items = ['🍎', '🍊', '🍌', '🥕', '🥬', '🍇', '🍓', '🥑', '🍅', '🥒']
    colors = ['#ff6b6b', '#4ecdc4', '#45b7d1', '#96ceb4', '#feca57', '#ff9ff3', '#54a0ff', '#5f27cd']
    
    # Create pattern
    for i in range(0, width, 100):
        for j in range(0, height, 100):
            # Add subtle circles
            circle_color = colors[(i//100 + j//100) % len(colors)]
            draw.ellipse([i+20, j+20, i+80, j+80], fill=circle_color + '20')
    
    return background

# Custom CSS with food theme
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Poppins:wght@300;400;600;700&display=swap');

.gradio-container {
    font-family: 'Poppins', sans-serif !important;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    min-height: 100vh;
}

.main-header {
    text-align: center;
    background: linear-gradient(135deg, #ff9a56, #ff6b95);
    color: white;
    padding: 3rem 2rem;
    border-radius: 20px;
    margin: 1rem;
    box-shadow: 0 10px 30px rgba(0, 0, 0, 0.3);
    position: relative;
    overflow: hidden;
}

.main-header::before {
    content: '🍎🥕🍊🥬🍇🍓';
    position: absolute;
    top: -10px;
    right: -10px;
    font-size: 2rem;
    opacity: 0.2;
    animation: float 3s ease-in-out infinite;
}

@keyframes float {
    0%, 100% { transform: translateY(0px) rotate(0deg); }
    50% { transform: translateY(-10px) rotate(5deg); }
}

.food-upload-area {
    background: linear-gradient(145deg, #ffffff, #f0f0f0);
    border: 3px dashed #ff9a56;
    border-radius: 20px;
    padding: 2rem;
    text-align: center;
    box-shadow: inset 0 2px 4px rgba(0, 0, 0, 0.1);
    transition: all 0.3s ease;
}

.food-upload-area:hover {
    border-color: #ff6b95;
    transform: translateY(-2px);
    box-shadow: 0 8px 25px rgba(255, 154, 86, 0.3);
}

.nutrition-panel {
    background: linear-gradient(145deg, #ffffff, #f8f9ff);
    border-radius: 15px;
    padding: 1.5rem;
    box-shadow: 0 5px 15px rgba(0, 0, 0, 0.1);
    border-left: 4px solid #ff9a56;
}

.btn-analyze {
    background: linear-gradient(135deg, #ff9a56, #ff6b95) !important;
    border: none !important;
    color: white !important;
    font-weight: 600 !important;
    padding: 12px 30px !important;
    border-radius: 25px !important;
    font-size: 16px !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 15px rgba(255, 154, 86, 0.4) !important;
}

.btn-analyze:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(255, 154, 86, 0.6) !important;
}

.food-facts {
    background: linear-gradient(145deg, #e8f5e8, #f0fff0);
    border-radius: 15px;
    padding: 1.5rem;
    margin-top: 1rem;
    border-left: 4px solid #4ecdc4;
}

.nutrition-grid {
    display: grid;
    grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
    gap: 1rem;
    margin-top: 1rem;
}

.nutrient-card {
    background: white;
    padding: 1rem;
    border-radius: 10px;
    text-align: center;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1);
    transition: transform 0.2s ease;
}

.nutrient-card:hover {
    transform: translateY(-3px);
}

/* Custom scrollbar */
::-webkit-scrollbar {
    width: 8px;
}

::-webkit-scrollbar-track {
    background: #f1f1f1;
    border-radius: 10px;
}

::-webkit-scrollbar-thumb {
    background: linear-gradient(135deg, #ff9a56, #ff6b95);
    border-radius: 10px;
}

::-webkit-scrollbar-thumb:hover {
    background: linear-gradient(135deg, #ff6b95, #ff9a56);
}
"""

def create_interface():
    """Create the Gradio interface"""
    with gr.Blocks(css=custom_css, title="🍎 AI Food Nutritionist", theme=gr.themes.Soft()) as interface:
        
        # Header with food theme
        gr.HTML("""
        <div class="main-header">
            <h1 style="font-size: 2.5rem; margin: 0; font-weight: 700;">🍎 AI Food Nutritionist</h1>
            <p style="font-size: 1.2rem; margin: 10px 0 0 0; opacity: 0.9;">Upload food images and discover nutritional insights with AI-powered analysis!</p>
        </div>
        """)
        
        with gr.Row():
            # Left panel - Image upload
            with gr.Column(scale=1):
                gr.HTML("<div style='text-align: center; padding: 1rem;'><h2 style='color: #333; margin-bottom: 1rem;'>📸 Upload Food Image</h2></div>")
                
                image_input = gr.Image(
                    label="Food Image",
                    type="pil",
                    height=400,
                    elem_classes=["food-upload-area"]
                )
                
                analyze_btn = gr.Button(
                    "🔍 Analyze Nutrition", 
                    variant="primary",
                    size="lg",
                    elem_classes=["btn-analyze"]
                )
                
                # Quick facts panel
                gr.HTML("""
                <div class="food-facts">
                    <h3 style="color: #2c3e50; margin-top: 0;">💡 Did you know?</h3>
                    <ul style="color: #555; line-height: 1.6;">
                        <li>🥕 Carrots contain beta-carotene for eye health</li>
                        <li>🍌 Bananas are rich in potassium for heart health</li>
                        <li>🥬 Leafy greens provide folate and iron</li>
                        <li>🍓 Berries are packed with antioxidants</li>
                    </ul>
                </div>
                """)
            
            # Right panel - Results
            with gr.Column(scale=2):
                gr.HTML("<div style='text-align: center; padding: 1rem;'><h2 style='color: #333; margin-bottom: 1rem;'>📊 Nutrition Analysis</h2></div>")
                
                with gr.Tabs():
                    with gr.TabItem("🍽️ Detailed Analysis"):
                        nutrition_output = gr.Textbox(
                            label="AI Nutrition Analysis",
                            lines=15,
                            interactive=False,
                            placeholder="Upload a food image and click 'Analyze Nutrition' to get detailed nutritional information...",
                            elem_classes=["nutrition-panel"]
                        )
                    
                    with gr.TabItem("🔍 Food Identification"):
                        identification_output = gr.Textbox(
                            label="Food Identification Details",
                            lines=10,
                            interactive=False,
                            placeholder="Food identification details will appear here...",
                            elem_classes=["nutrition-panel"]
                        )
                    
                    with gr.TabItem("📈 Quick Stats"):
                        stats_output = gr.Textbox(
                            label="Quick Nutritional Breakdown",
                            lines=10,
                            interactive=False,
                            placeholder="Quick nutritional statistics will appear here...",
                            elem_classes=["nutrition-panel"]
                        )
        
        # Example images section
        gr.HTML("""
        <div style="margin-top: 2rem; text-align: center; padding: 2rem; background: rgba(255, 255, 255, 0.1); border-radius: 15px; backdrop-filter: blur(10px);">
            <h3 style="color: white; margin-bottom: 1rem;">🌟 Try These Examples</h3>
            <p style="color: white; opacity: 0.9;">Upload images of fruits, vegetables, nuts, or other foods to get instant nutritional analysis!</p>
            <div style="margin-top: 1rem; font-size: 2rem;">
                🍎 🥕 🍌 🥬 🍇 🍓 🥑 🍅 🥒 🌽
            </div>
        </div>
        """)
        
        # Processing status
        status_display = gr.HTML(visible=False)
        
        # Event handler
        def analyze_food_wrapper(image):
            if image is None:
                return "Please upload an image first! 📸", "", ""
            
            # Show processing status
            status_html = """
            <div style="text-align: center; padding: 1rem; background: #e3f2fd; border-radius: 10px; margin: 1rem;">
                <h3 style="color: #1976d2;">🔄 Processing your food image...</h3>
                <p>AI is analyzing the nutritional content. Please wait...</p>
            </div>
            """
            
            try:
                nutrition_info, breakdown, identification = food_app.process_food_image(image)
                return nutrition_info, breakdown, identification
            except Exception as e:
                error_msg = f"❌ Error analyzing image: {str(e)}"
                return error_msg, "", ""
        
        analyze_btn.click(
            fn=analyze_food_wrapper,
            inputs=[image_input],
            outputs=[nutrition_output, stats_output, identification_output]
        )
    
    return interface

# Launch the application
if __name__ == "__main__":
    interface = create_interface()
    interface.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )