"""
AI Nutrition Tracker with Image Recognition
Author: RSK World (https://rskworld.in)
Founded by: Molla Samser
Designer & Tester: Rima Khatun
Contact: help@rskworld.in, +91 93305 39277
Year: 2026
"""

import base64
import json
import requests
from datetime import datetime, timedelta
from typing import Dict, List, Any, Optional

class NutritionAI:
    """Advanced AI-powered nutrition tracking with image recognition"""
    
    def __init__(self):
        self.food_database = self._load_food_database()
        self.nutrition_goals = self._load_nutrition_goals()
        self.meal_patterns = {}
        
    def _load_food_database(self) -> Dict:
        """Load comprehensive food nutrition database"""
        return {
            # Fruits
            "apple": {
                "calories": 52, "protein": 0.3, "carbs": 14, "fat": 0.2, "fiber": 2.4,
                "vitamins": {"C": 8, "A": 54}, "category": "fruit", "glycemic_index": 38
            },
            "banana": {
                "calories": 89, "protein": 1.1, "carbs": 23, "fat": 0.3, "fiber": 2.6,
                "vitamins": {"C": 15, "B6": 20}, "category": "fruit", "glycemic_index": 51
            },
            "orange": {
                "calories": 47, "protein": 0.9, "carbs": 12, "fat": 0.1, "fiber": 2.4,
                "vitamins": {"C": 88, "A": 4}, "category": "fruit", "glycemic_index": 43
            },
            
            # Vegetables
            "broccoli": {
                "calories": 34, "protein": 2.8, "carbs": 7, "fat": 0.4, "fiber": 2.6,
                "vitamins": {"C": 135, "K": 116}, "category": "vegetable", "glycemic_index": 10
            },
            "spinach": {
                "calories": 23, "protein": 2.9, "carbs": 3.6, "fat": 0.4, "fiber": 2.2,
                "vitamins": {"K": 402, "A": 105}, "category": "vegetable", "glycemic_index": 15
            },
            "carrot": {
                "calories": 41, "protein": 0.9, "carbs": 10, "fat": 0.2, "fiber": 2.8,
                "vitamins": {"A": 835, "K": 13}, "category": "vegetable", "glycemic_index": 47
            },
            
            # Proteins
            "chicken_breast": {
                "calories": 165, "protein": 31, "carbs": 0, "fat": 3.6, "fiber": 0,
                "vitamins": {"B6": 30, "B12": 6}, "category": "protein", "glycemic_index": 0
            },
            "salmon": {
                "calories": 208, "protein": 20, "carbs": 0, "fat": 13, "fiber": 0,
                "vitamins": {"D": 236, "B12": 209}, "category": "protein", "glycemic_index": 0
            },
            "eggs": {
                "calories": 155, "protein": 13, "carbs": 1.1, "fat": 11, "fiber": 0,
                "vitamins": {"B12": 46, "D": 21}, "category": "protein", "glycemic_index": 42
            },
            
            # Grains
            "rice_white": {
                "calories": 130, "protein": 2.7, "carbs": 28, "fat": 0.3, "fiber": 0.4,
                "vitamins": {"B1": 8}, "category": "grain", "glycemic_index": 73
            },
            "rice_brown": {
                "calories": 111, "protein": 2.6, "carbs": 23, "fat": 0.9, "fiber": 1.8,
                "vitamins": {"B1": 12}, "category": "grain", "glycemic_index": 50
            },
            "quinoa": {
                "calories": 120, "protein": 4.4, "carbs": 21, "fat": 1.9, "fiber": 2.8,
                "vitamins": {"B1": 13, "E": 6}, "category": "grain", "glycemic_index": 35
            },
            
            # Dairy
            "milk": {
                "calories": 42, "protein": 3.4, "carbs": 5, "fat": 1, "fiber": 0,
                "vitamins": {"B12": 18, "D": 51}, "category": "dairy", "glycemic_index": 46
            },
            "yogurt": {
                "calories": 59, "protein": 10, "carbs": 3.6, "fat": 0.4, "fiber": 0,
                "vitamins": {"B12": 17}, "category": "dairy", "glycemic_index": 47
            },
            
            # Nuts & Seeds
            "almonds": {
                "calories": 579, "protein": 21, "carbs": 22, "fat": 50, "fiber": 12.5,
                "vitamins": {"E": 173}, "category": "nuts", "glycemic_index": 0
            },
            "chia_seeds": {
                "calories": 486, "protein": 17, "carbs": 42, "fat": 31, "fiber": 34.4,
                "vitamins": {"E": 14}, "category": "seeds", "glycemic_index": 1
            }
        }
    
    def _load_nutrition_goals(self) -> Dict:
        """Load nutrition goals based on user profiles"""
        return {
            "weight_loss": {
                "calories": 1800,
                "protein": 150,  # grams
                "carbs": 180,
                "fat": 60,
                "fiber": 25
            },
            "muscle_gain": {
                "calories": 2800,
                "protein": 180,
                "carbs": 280,
                "fat": 90,
                "fiber": 30
            },
            "maintenance": {
                "calories": 2200,
                "protein": 120,
                "carbs": 220,
                "fat": 73,
                "fiber": 25
            },
            "endurance": {
                "calories": 2500,
                "protein": 130,
                "carbs": 312,
                "fat": 83,
                "fiber": 28
            }
        }
    
    async def analyze_food_image(self, image_data: str, user_profile: Dict) -> Dict:
        """Analyze food image using AI vision"""
        try:
            # Prepare image for AI analysis
            image_base64 = image_data.split(',')[1] if ',' in image_data else image_data
            
            # Use OpenAI Vision API for food recognition
            response = await self._call_vision_api(image_base64)
            
            # Parse the response and identify foods
            identified_foods = self._parse_food_response(response)
            
            # Calculate nutrition information
            nutrition_analysis = self._calculate_meal_nutrition(identified_foods, user_profile)
            
            # Generate recommendations
            recommendations = self._generate_nutrition_recommendations(nutrition_analysis, user_profile)
            
            return {
                "success": True,
                "identified_foods": identified_foods,
                "nutrition_analysis": nutrition_analysis,
                "recommendations": recommendations,
                "timestamp": datetime.now().isoformat()
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "fallback_analysis": self._fallback_food_analysis()
            }
    
    async def _call_vision_api(self, image_base64: str) -> str:
        """Call vision API for food recognition"""
        # This would integrate with OpenAI Vision API or similar
        # For now, we'll simulate the response
        prompt = """
        Analyze this food image and identify all visible food items. 
        For each food item, provide:
        1. Name of the food
        2. Estimated portion size (in grams or common measurements)
        3. Confidence level (0-100)
        4. Preparation method (grilled, fried, raw, etc.)
        
        Format as JSON array.
        """
        
        # Simulated response - replace with actual API call
        return json.dumps([
            {
                "name": "grilled_chicken_breast",
                "portion": "150g",
                "confidence": 95,
                "preparation": "grilled"
            },
            {
                "name": "broccoli",
                "portion": "100g",
                "confidence": 88,
                "preparation": "steamed"
            },
            {
                "name": "rice_brown",
                "portion": "200g",
                "confidence": 92,
                "preparation": "boiled"
            }
        ])
    
    def _parse_food_response(self, response: str) -> List[Dict]:
        """Parse AI vision response to identify foods"""
        try:
            foods_data = json.loads(response)
            identified_foods = []
            
            for food in foods_data:
                # Map to our database
                food_name = self._normalize_food_name(food["name"])
                if food_name in self.food_database:
                    nutrition = self.food_database[food_name]
                    portion_grams = self._estimate_portion_grams(food["portion"])
                    
                    identified_foods.append({
                        "name": food_name,
                        "display_name": food["name"].replace("_", " ").title(),
                        "portion": food["portion"],
                        "portion_grams": portion_grams,
                        "confidence": food["confidence"],
                        "preparation": food["preparation"],
                        "nutrition": self._scale_nutrition(nutrition, portion_grams, 100)
                    })
            
            return identified_foods
            
        except Exception as e:
            print(f"Error parsing food response: {e}")
            return []
    
    def _normalize_food_name(self, food_name: str) -> str:
        """Normalize food name to match database"""
        # Remove preparation methods and normalize
        normalized = food_name.lower().replace("grilled_", "").replace("fried_", "").replace("baked_", "")
        
        # Map common variations
        name_mapping = {
            "chicken": "chicken_breast",
            "rice": "rice_white" if "white" in food_name else "rice_brown",
            "bread": "whole_wheat_bread"
        }
        
        return name_mapping.get(normalized, normalized)
    
    def _estimate_portion_grams(self, portion: str) -> float:
        """Estimate portion size in grams"""
        portion = portion.lower().strip()
        
        # Common portion sizes
        portion_sizes = {
            "small": 100,
            "medium": 150,
            "large": 200,
            "cup": 240,
            "half_cup": 120,
            "tablespoon": 15,
            "teaspoon": 5,
            "slice": 30,
            "piece": 50
        }
        
        # Extract numbers from portion string
        import re
        numbers = re.findall(r'\d+', portion)
        if numbers:
            return float(numbers[0])
        
        # Check for common terms
        for term, grams in portion_sizes.items():
            if term in portion:
                return grams
        
        return 100  # Default portion
    
    def _scale_nutrition(self, base_nutrition: Dict, portion_grams: float, base_grams: float) -> Dict:
        """Scale nutrition based on portion size"""
        scale_factor = portion_grams / base_grams
        scaled = {}
        
        for key, value in base_nutrition.items():
            if isinstance(value, (int, float)):
                scaled[key] = round(value * scale_factor, 2)
            elif isinstance(value, dict):
                scaled[key] = {k: round(v * scale_factor, 2) for k, v in value.items()}
            else:
                scaled[key] = value
        
        return scaled
    
    def _calculate_meal_nutrition(self, foods: List[Dict], user_profile: Dict) -> Dict:
        """Calculate total nutrition for the meal"""
        total_nutrition = {
            "calories": 0,
            "protein": 0,
            "carbs": 0,
            "fat": 0,
            "fiber": 0,
            "vitamins": {},
            "categories": {}
        }
        
        for food in foods:
            nutrition = food["nutrition"]
            total_nutrition["calories"] += nutrition.get("calories", 0)
            total_nutrition["protein"] += nutrition.get("protein", 0)
            total_nutrition["carbs"] += nutrition.get("carbs", 0)
            total_nutrition["fat"] += nutrition.get("fat", 0)
            total_nutrition["fiber"] += nutrition.get("fiber", 0)
            
            # Sum vitamins
            for vitamin, amount in nutrition.get("vitamins", {}).items():
                total_nutrition["vitamins"][vitamin] = total_nutrition["vitamins"].get(vitamin, 0) + amount
            
            # Track categories
            category = nutrition.get("category", "unknown")
            total_nutrition["categories"][category] = total_nutrition["categories"].get(category, 0) + nutrition.get("calories", 0)
        
        # Calculate macros percentages
        total_macros = total_nutrition["protein"] * 4 + total_nutrition["carbs"] * 4 + total_nutrition["fat"] * 9
        if total_macros > 0:
            total_nutrition["macro_percentages"] = {
                "protein": round((total_nutrition["protein"] * 4 / total_macros) * 100, 1),
                "carbs": round((total_nutrition["carbs"] * 4 / total_macros) * 100, 1),
                "fat": round((total_nutrition["fat"] * 9 / total_macros) * 100, 1)
            }
        
        # Compare to goals
        goal = self.nutrition_goals.get(user_profile.get("goal", "maintenance"), self.nutrition_goals["maintenance"])
        total_nutrition["goal_comparison"] = {
            "calories": {
                "current": total_nutrition["calories"],
                "goal": goal["calories"],
                "percentage": round((total_nutrition["calories"] / goal["calories"]) * 100, 1)
            },
            "protein": {
                "current": total_nutrition["protein"],
                "goal": goal["protein"],
                "percentage": round((total_nutrition["protein"] / goal["protein"]) * 100, 1)
            }
        }
        
        return total_nutrition
    
    def _generate_nutrition_recommendations(self, nutrition: Dict, user_profile: Dict) -> List[str]:
        """Generate personalized nutrition recommendations"""
        recommendations = []
        goal = user_profile.get("goal", "maintenance")
        
        # Calorie recommendations
        calorie_percentage = nutrition["goal_comparison"]["calories"]["percentage"]
        if calorie_percentage > 120:
            recommendations.append("This meal is high in calories. Consider reducing portion sizes for your goal.")
        elif calorie_percentage < 80:
            recommendations.append("This meal is relatively low in calories. You could add healthy fats or complex carbs.")
        
        # Protein recommendations
        protein_percentage = nutrition["goal_comparison"]["protein"]["percentage"]
        if protein_percentage < 80 and goal in ["muscle_gain", "weight_loss"]:
            recommendations.append("Consider adding more protein to support muscle maintenance and growth.")
        
        # Balance recommendations
        macros = nutrition.get("macro_percentages", {})
        if macros.get("protein", 0) < 20:
            recommendations.append("Increase protein intake for better satiety and muscle support.")
        if macros.get("carbs", 0) > 60:
            recommendations.append("Consider reducing carbohydrates and increasing healthy fats.")
        
        # Fiber recommendations
        if nutrition.get("fiber", 0) < 5:
            recommendations.append("Add more vegetables or whole grains to increase fiber intake.")
        
        # Category balance
        categories = nutrition.get("categories", {})
        if categories.get("vegetable", 0) < 100:
            recommendations.append("Include more non-starchy vegetables for micronutrients and fiber.")
        
        # Goal-specific recommendations
        if goal == "muscle_gain":
            if protein_percentage < 100:
                recommendations.append("For muscle gain, aim for higher protein intake (1.6-2.2g per kg bodyweight).")
        elif goal == "weight_loss":
            if calorie_percentage > 100:
                recommendations.append("For weight loss, create a moderate calorie deficit through portion control.")
        
        return recommendations
    
    def _fallback_food_analysis(self) -> Dict:
        """Fallback analysis when AI vision fails"""
        return {
            "message": "Unable to analyze image. Please try again with better lighting or manually log your meal.",
            "suggestions": [
                "Ensure food is well-lit and clearly visible",
                "Take photo from above (top-down view)",
                "Avoid shadows and glare",
                "Use manual entry as backup"
            ]
        }
    
    def track_daily_nutrition(self, user_id: str, date: str) -> Dict:
        """Track and analyze daily nutrition"""
        # This would integrate with database to track daily intake
        # For now, return sample analysis
        return {
            "date": date,
            "total_calories": 1850,
            "total_protein": 142,
            "total_carbs": 198,
            "total_fat": 65,
            "meals_logged": 3,
            "water_intake": 2.1,  # liters
            "goal_progress": {
                "calories": 84,
                "protein": 95,
                "water": 84
            },
            "recommendations": [
                "Increase water intake to reach 2.5L daily goal",
                "Add a protein-rich snack to meet protein goals",
                "Consider adding more vegetables to dinner"
            ]
        }
    
    def generate_meal_plan(self, user_profile: Dict, preferences: Dict) -> Dict:
        """Generate personalized meal plan"""
        goal = user_profile.get("goal", "maintenance")
        calories = user_profile.get("daily_calories", 2200)
        dietary_restrictions = preferences.get("restrictions", [])
        meal_frequency = preferences.get("meal_frequency", 3)
        
        # Calculate calories per meal
        calories_per_meal = calories // meal_frequency
        
        meal_plan = {
            "goal": goal,
            "daily_calories": calories,
            "meals": [],
            "shopping_list": [],
            "prep_tips": []
        }
        
        # Generate meals for each meal type
        meal_types = ["breakfast", "lunch", "dinner", "snack"]
        for meal_type in meal_types[:meal_frequency]:
            meal = self._generate_meal(meal_type, calories_per_meal, dietary_restrictions)
            meal_plan["meals"].append(meal)
        
        # Generate shopping list
        meal_plan["shopping_list"] = self._generate_shopping_list(meal_plan["meals"])
        
        # Add prep tips
        meal_plan["prep_tips"] = [
            "Prep vegetables in advance for quick meal assembly",
            "Cook proteins in batches and store for the week",
            "Use portion control containers for accurate serving sizes",
            "Keep healthy snacks readily available"
        ]
        
        return meal_plan
    
    def _generate_meal(self, meal_type: str, calories: int, restrictions: List[str]) -> Dict:
        """Generate a single meal"""
        meal_templates = {
            "breakfast": [
                {"name": "Protein Oatmeal", "foods": ["oats", "eggs", "berries", "almonds"]},
                {"name": "Greek Yogurt Parfait", "foods": ["yogurt", "granola", "fruits", "chia_seeds"]},
                {"name": "Scrambled Eggs with Toast", "foods": ["eggs", "whole_wheat_bread", "spinach", "milk"]}
            ],
            "lunch": [
                {"name": "Grilled Chicken Salad", "foods": ["chicken_breast", "mixed_greens", "vegetables", "olive_oil"]},
                {"name": "Quinoa Buddha Bowl", "foods": ["quinoa", "chickpeas", "vegetables", "tahini"]},
                {"name": "Turkey Wrap", "foods": ["turkey", "whole_wheat_tortilla", "vegetables", "hummus"]}
            ],
            "dinner": [
                {"name": "Salmon with Vegetables", "foods": ["salmon", "broccoli", "rice_brown", "olive_oil"]},
                {"name": "Lean Beef Stir-fry", "foods": ["lean_beef", "vegetables", "rice_brown", "soy_sauce"]},
                {"name": "Chicken Vegetable Curry", "foods": ["chicken_breast", "vegetables", "coconut_milk", "spices"]}
            ],
            "snack": [
                {"name": "Protein Smoothie", "foods": ["protein_powder", "banana", "almonds", "milk"]},
                {"name": "Apple with Almond Butter", "foods": ["apple", "almond_butter", "cinnamon"]},
                {"name": "Greek Yogurt with Berries", "foods": ["yogurt", "berries", "honey"]}
            ]
        }
        
        templates = meal_templates.get(meal_type, meal_templates["lunch"])
        selected_template = templates[0]  # In production, would rotate based on variety
        
        meal = {
            "type": meal_type,
            "name": selected_template["name"],
            "foods": [],
            "total_calories": 0,
            "total_protein": 0,
            "instructions": self._get_meal_instructions(selected_template["name"])
        }
        
        # Add foods with nutrition
        for food_name in selected_template["foods"]:
            if food_name in self.food_database:
                nutrition = self.food_database[food_name]
                meal["foods"].append({
                    "name": food_name,
                    "display_name": food_name.replace("_", " ").title(),
                    "nutrition": nutrition
                })
                meal["total_calories"] += nutrition.get("calories", 0)
                meal["total_protein"] += nutrition.get("protein", 0)
        
        return meal
    
    def _get_meal_instructions(self, meal_name: str) -> str:
        """Get cooking instructions for meal"""
        instructions = {
            "Protein Oatmeal": "Cook oats with water or milk, stir in protein powder, top with berries and almonds",
            "Grilled Chicken Salad": "Grill chicken breast, slice and serve over mixed greens with vegetables",
            "Salmon with Vegetables": "Bake salmon at 400°F for 15-20 minutes, serve with steamed vegetables",
            "Protein Smoothie": "Blend all ingredients until smooth, add ice if desired"
        }
        
        return instructions.get(meal_name, "Follow standard cooking instructions for best results")
    
    def _generate_shopping_list(self, meals: List[Dict]) -> List[Dict]:
        """Generate shopping list from meal plan"""
        shopping_list = {}
        
        for meal in meals:
            for food in meal["foods"]:
                food_name = food["name"]
                shopping_list[food_name] = shopping_list.get(food_name, 0) + 1
        
        return [
            {"item": food.replace("_", " ").title(), "quantity": quantity}
            for food, quantity in shopping_list.items()
        ]
    
    def analyze_nutrition_trends(self, user_id: str, days: int = 30) -> Dict:
        """Analyze nutrition trends over time"""
        # This would query database for historical data
        # Return sample analysis for now
        return {
            "period": f"Last {days} days",
            "average_daily_calories": 2150,
            "average_daily_protein": 135,
            "goal_consistency": 78,  # percentage
            "most_consumed_foods": [
                {"food": "chicken_breast", "frequency": 12},
                {"food": "rice_brown", "frequency": 10},
                {"food": "broccoli", "frequency": 8}
            ],
            "improvement_areas": [
                "Increase vegetable variety",
                "Add more healthy fats",
                "Maintain consistent protein intake"
            ],
            "achievements": [
                {"name": "7-Day Streak", "description": "Logged meals for 7 consecutive days"},
                {"name": "Protein Goal", "description": "Met protein goals 5 times this week"}
            ]
        }