agents/specialized/nutrition_agent.py

"""
Nutrition Agent - Specialized agent for nutrition advice
"""

from config.settings import client, MODEL
from modules.nutrition import NutritionAdvisor
from health_data import HealthContext
from personalization import PersonalizationEngine
from rag.rag_integration import get_rag_integration
from agents.core.base_agent import BaseAgent
from agents.core.context_analyzer import ContextAnalyzer
from agents.core.response_validator import ResponseValidator
from typing import Dict, Any, List, Optional
from datetime import datetime
import re

class NutritionAgent(BaseAgent):
    def __init__(self, memory=None):
        super().__init__(memory)
        self.advisor = NutritionAdvisor()
        self.health_context = None
        self.personalization = None
        self.rag = get_rag_integration()
        
        # Configure handoff triggers for nutrition agent
        self.handoff_triggers = {
            'exercise_agent': ['tập', 'gym', 'cardio', 'yoga', 'chạy bộ', 'thể dục', 'vận động'],
            'symptom_agent': ['đau bụng', 'buồn nôn', 'tiêu chảy', 'dị ứng', 'ngộ độc'],
            'mental_health_agent': ['stress', 'lo âu', 'mất ngủ', 'ăn không ngon'],
            'general_health_agent': ['khám', 'xét nghiệm', 'bác sĩ']
        }
        self.system_prompt = """Bạn là chuyên gia dinh dưỡng chuyên nghiệp.

🥗 CHUYÊN MÔN:
- Tư vấn dinh dưỡng cá nhân hóa dựa trên BMI, tuổi, giới tính, mục tiêu
- Tính toán calo, macro (protein/carb/fat)
- Gợi ý thực đơn phù hợp
- Tư vấn thực phẩm bổ sung
- Hướng dẫn ăn uống cho các bệnh lý (tiểu đường, huyết áp, tim mạch...)

🎯 CÁCH TƯ VẤN:

1. **KIỂM TRA THÔNG TIN TRƯỚC KHI HỎI:**
   - ĐỌC KỸ chat history - user có thể đã cung cấp thông tin rồi!
   - Nếu user đã nói "tôi 25 tuổi, nam, 70kg, 175cm" → ĐỪNG HỎI LẠI!
   - Chỉ hỏi thông tin THỰC SỰ còn thiếu
   - Nếu đã đủ (tuổi, giới tính, cân nặng, chiều cao) → ĐƯA KHUYẾN NGHỊ NGAY!

2. **ƯU TIÊN THÔNG TIN:**
   - Câu 1: Mục tiêu (giảm cân/tăng cân/duy trì?)
   - Câu 2: Cân nặng, chiều cao (để tính BMI)
   - Câu 3: Mức độ hoạt động (ít/vừa/nhiều)
   - Câu 4 (nếu cần): Bệnh nền, dị ứng

3. **KHI USER KHÔNG MUỐN CUNG CẤP:**
   - User nói "không biết", "không muốn nói", "tư vấn chung thôi"
   - → DỪNG hỏi, đưa khuyến nghị chung
   - Dựa trên thông tin ĐÃ CÓ để tư vấn

4. **ĐƯA KHUYẾN NGHỊ:**
   - Nếu có đủ thông tin: Tính calo, macro cụ thể
   - Nếu thiếu thông tin: Đưa khuyến nghị chung (400g rau củ, protein đủ, etc.)
   - Gợi ý thực đơn mẫu
   - KHÔNG hỏi thêm nữa

⚠️ AN TOÀN:
- Luôn khuyên gặp bác sĩ dinh dưỡng cho các vấn đề phức tạp
- Cảnh báo về các chế độ ăn kiêng cực đoan
- Lưu ý về dị ứng, bệnh nền

💬 PHONG CÁCH:
- Chuyên nghiệp, rõ ràng, súc tích
- Dùng "tôi" để thể hiện tính chuyên môn
- KHÔNG dùng emoji
- Đưa ra con số cụ thể khi có thể
- Thực tế, không lý thuyết suông
- TỰ NHIÊN, MẠCH LẠC - không lặp lại ý, không copy-paste câu từ context khác
- Nếu hỏi thông tin → Hỏi NGẮN GỌN, TRỰC TIẾP
- KHÔNG dùng câu như "Bạn thử làm theo xem có đỡ không" (đây là câu của bác sĩ chữa bệnh!)"""

    def set_health_context(self, health_context: HealthContext):
        """Inject health context and initialize personalization engine"""
        self.health_context = health_context
        self.personalization = PersonalizationEngine(health_context)

    def handle(self, parameters, chat_history=None):
        """
        Handle nutrition request using LLM for natural conversation

        Args:
            parameters (dict): {
                "user_query": str,
                "user_data": dict (optional)
            }
            chat_history (list): Conversation history

        Returns:
            str: Response message
        """
        user_query = parameters.get("user_query", "")
        user_data = parameters.get("user_data", {})
        
        # Extract and save user info from current message immediately
        self.extract_and_save_user_info(user_query)
        
        # Update memory from chat history
        if chat_history:
            self.update_memory_from_history(chat_history)
        
        # Check if we should hand off to another agent
        if self.should_handoff(user_query, chat_history):
            next_agent = self.suggest_next_agent(user_query)
            if next_agent:
                # Save current nutrition data for next agent
                self.save_agent_data('last_nutrition_advice', {
                    'query': user_query,
                    'user_profile': self.get_user_profile(),
                    'timestamp': datetime.now().isoformat()
                })
                
                # Create handoff message with context
                context = self._generate_nutrition_summary()
                return self.create_handoff_message(next_agent, context, user_query)

        # Use health context if available
        if self.health_context:
            profile = self.health_context.get_user_profile()
            user_data = {
                'age': profile.age,
                'gender': profile.gender,
                'weight': profile.weight,
                'height': profile.height,
                'activity_level': profile.activity_level,
                'health_conditions': profile.health_conditions,
                'dietary_restrictions': profile.dietary_restrictions
            }
        # Extract user data from chat history if not provided
        elif not user_data and chat_history:
            user_data = self._extract_user_data_from_history(chat_history)
            # Save extracted data to shared memory for other agents
            for key, value in user_data.items():
                if value is not None:
                    self.update_user_profile(key, value)

        # Check if user needs personalized advice (BMI, calories, meal plan)
        needs_personalization = self._needs_personalized_advice(user_query, chat_history)

        if needs_personalization:
            # Check if we have enough data
            missing_fields = self._check_missing_data(user_data)

            if missing_fields:
                return self._ask_for_missing_data(missing_fields, user_data, user_query)

            # Generate personalized nutrition advice
            try:
                result = self.advisor.generate_nutrition_advice(user_data)

                # Adapt recommendations using personalization engine
                if self.personalization:
                    adapted_result = self.personalization.adapt_nutrition_plan(result)
                else:
                    adapted_result = result

                response = self._format_nutrition_response(adapted_result, user_data)

                # Persist data to health context
                if self.health_context:
                    self.health_context.add_health_record('nutrition', {
                        'query': user_query,
                        'advice': response,
                        'user_data': user_data,
                        'timestamp': datetime.now().isoformat()
                    })

                return response
            except Exception as e:
                return self._handle_error(e, user_query)
        else:
            # General nutrition question - use LLM directly
            response = self._handle_general_nutrition_query(user_query, chat_history)

            # Persist general query
            if self.health_context:
                self.health_context.add_health_record('nutrition', {
                    'query': user_query,
                    'response': response,
                    'type': 'general',
                    'timestamp': datetime.now().isoformat()
                })

            return response
    
    def _extract_user_data_from_history(self, chat_history):
        """Extract user data from conversation history"""
        user_data = {
            'age': None,
            'gender': None,
            'weight': None,
            'height': None,
            'goal': 'maintenance',
            'activity_level': 'moderate',
            'dietary_restrictions': [],
            'health_conditions': []
        }
        
        all_messages = " ".join([msg[0] for msg in chat_history if msg[0]])
        
        # Extract age
        age_match = re.search(r'(\d+)\s*tuổi|tuổi\s*(\d+)|tôi\s*(\d+)', all_messages.lower())
        if age_match:
            user_data['age'] = int([g for g in age_match.groups() if g][0])
        
        # Extract gender
        if re.search(r'\bnam\b|male|đàn ông', all_messages.lower()):
            user_data['gender'] = 'male'
        elif re.search(r'\bnữ\b|female|đàn bà', all_messages.lower()):
            user_data['gender'] = 'female'
        
        # Extract weight - improved patterns
        weight_match = re.search(r'(?:nặng|cân|weight)?\s*(\d+(?:\.\d+)?)\s*kg|(\d+(?:\.\d+)?)\s*kg', all_messages.lower())
        if weight_match:
            user_data['weight'] = float([g for g in weight_match.groups() if g][0])
        
        # Extract height - improved patterns
        height_cm_match = re.search(r'(?:cao|chiều cao|height)?\s*(\d+(?:\.\d+)?)\s*cm', all_messages.lower())
        if height_cm_match:
            user_data['height'] = float(height_cm_match.group(1))
        else:
            height_m_match = re.search(r'(?:cao|chiều cao|height)?\s*(\d+\.?\d*)\s*m\b', all_messages.lower())
            if height_m_match:
                height = float(height_m_match.group(1))
                if height < 3:  # Convert meters to cm
                    height = height * 100
                user_data['height'] = height
        
        # Extract goal
        if re.search(r'giảm cân|weight loss', all_messages.lower()):
            user_data['goal'] = 'weight_loss'
        elif re.search(r'tăng cân|weight gain', all_messages.lower()):
            user_data['goal'] = 'weight_gain'
        elif re.search(r'tập gym|muscle|cơ bắp', all_messages.lower()):
            user_data['goal'] = 'muscle_building'
        
        return user_data
    
    def _needs_personalized_advice(self, user_query, chat_history):
        """
        Determine if user needs personalized advice (BMI, calories, meal plan)
        or just general nutrition info
        """
        # Keywords that indicate need for personalization
        personalization_keywords = [
            'giảm cân', 'tăng cân', 'bmi', 'calo', 'calorie',
            'thực đơn', 'meal plan', 'chế độ ăn cá nhân',
            'tôi nên ăn gì', 'tư vấn cho tôi', 'phù hợp với tôi'
        ]
        
        query_lower = user_query.lower()
        
        # Check if user explicitly asks for personalized advice
        if any(kw in query_lower for kw in personalization_keywords):
            return True
        
        # Check chat history - if user already provided personal info
        if chat_history:
            all_messages = " ".join([msg[0] for msg in chat_history if msg[0]]).lower()
            if any(kw in all_messages for kw in personalization_keywords):
                return True
        
        # Default: general question
        return False
    
    def _check_missing_data(self, user_data):
        """Check what data is missing - check shared memory first"""
        required = ['age', 'gender', 'weight', 'height']
        
        # Check shared memory for missing fields
        profile = self.get_user_profile()
        for field in required:
            if not user_data.get(field) and profile.get(field):
                user_data[field] = profile[field]
        
        return [field for field in required if not user_data.get(field)]
    
    def _ask_for_missing_data(self, missing_fields, current_data, user_query):
        """Ask for missing data"""
        questions = {
            'age': "bạn bao nhiêu tuổi",
            'gender': "bạn là nam hay nữ",
            'weight': "bạn nặng bao nhiêu kg",
            'height': "bạn cao bao nhiêu cm"
        }
        
        # Build friendly question
        q_list = [questions[f] for f in missing_fields]
        
        if len(q_list) == 1:
            question = q_list[0]
        elif len(q_list) == 2:
            question = f"{q_list[0]} và {q_list[1]}"
        else:
            question = ", ".join(q_list[:-1]) + f" và {q_list[-1]}"
        
        return f"""🥗 **Để tư vấn dinh dưỡng chính xác, mình cần biết thêm:**

Cho mình biết {question} nhé?

💡 **Ví dụ:** "Tôi 25 tuổi, nam, nặng 70kg, cao 175cm"

Sau khi có đủ thông tin, mình sẽ tính BMI và đưa ra lời khuyên dinh dưỡng cá nhân hóa cho bạn! 😊"""
    
    def _format_nutrition_response(self, result, user_data):
        """Format nutrition advice into friendly response"""
        bmi_info = result['bmi_analysis']
        targets = result['daily_targets']
        meals = result['meal_suggestions']
        supplements = result['supplement_recommendations']
        
        response = f"""🥗 **Tư Vấn Dinh Dưỡng Cá Nhân Hóa**

👤 **Thông tin của bạn:**
- {user_data['age']} tuổi, {user_data['gender']}, {user_data['weight']}kg, {user_data['height']}cm

📊 **Phân tích BMI:**
- BMI: **{bmi_info['bmi']}** ({bmi_info['category']})
- Lời khuyên: {bmi_info['advice']}

🎯 **Mục tiêu hàng ngày:**
- 🔥 Calo: **{targets['daily_calories']} kcal**
- 🥩 Protein: **{targets['protein']}**
- 🍚 Carb: **{targets['carbs']}**
- 🥑 Chất béo: **{targets['fats']}**
- 💧 Nước: **{targets['water']}**

🍽️ **Gợi ý thực đơn:**

**Sáng:**
- {meals['breakfast'][0]}
- {meals['breakfast'][1]}

**Trưa:**
- {meals['lunch'][0]}
- {meals['lunch'][1]}

**Tối:**
- {meals['dinner'][0]}
- {meals['dinner'][1]}

**Snack:**
- {meals['snacks'][0]}
- {meals['snacks'][1]}
"""
        
        if supplements:
            response += f"\n💊 **Thực phẩm bổ sung gợi ý:**\n"
            response += "\n".join([f"- {s}" for s in supplements[:4]])
        
        response += f"""

🤖 **Lời khuyên chuyên gia:**
{result['personalized_advice'][:600]}...

---

⚠️ *Đây là tư vấn tham khảo. Với các vấn đề phức tạp, hãy gặp bác sĩ dinh dưỡng nhé!*

💬 Bạn có câu hỏi gì về chế độ ăn này không? Hoặc muốn mình điều chỉnh gì không? 😊"""
        
        return response
    
    def _build_nutrition_context_instruction(self, user_query, chat_history):
        """
        Build context instruction for nutrition queries
        """
        # Check if user is answering comparison self-assessment
        if chat_history and len(chat_history) > 0:
            last_bot_msg = chat_history[-1][1] if len(chat_history[-1]) > 1 else ""
            if "TỰ KIỂM TRA" in last_bot_msg or "Bạn trả lời" in last_bot_msg:
                return """\n\nPHASE: PHÂN TÍCH LỰA CHỌN DINH DƯỠNG
User vừa trả lời các câu hỏi. Phân tích:

1. NHẬN DIỆN PHÙ HỢP (dựa vào RAG):
   - Đọc kỹ mục tiêu, lifestyle, sở thích
   - So sánh với đặc điểm của từng lựa chọn
   - Đưa ra lựa chọn PHÙ HỢP NHẤT

2. GIẢI THÍCH:
   - Vì sao lựa chọn này phù hợp
   - Lợi ích cụ thể cho user
   - Lưu ý khi thực hiện

3. HƯỚNG DẪN BẮT ĐẦU:
   - Cách bắt đầu cụ thể
   - Thực đơn mẫu (nếu cần)
   - Tips để duy trì

4. Kết thúc: "Bạn cần hướng dẫn chi tiết hơn không?"
KHÔNG nói "Dựa trên thông tin"."""
        
        # Check if asking comparison question
        if any(phrase in user_query.lower() for phrase in [
            'nên ăn', 'hay', 'hoặc', 'khác nhau thế nào',
            'chọn', 'so sánh', 'tốt hơn'
        ]):
            return """\n\nPHASE: SO SÁNH DINH DƯỠNG (GENERIC)
User muốn so sánh các lựa chọn dinh dưỡng. Sử dụng RAG để:

1. XÁC ĐỊNH các lựa chọn (từ user query):
   - Trích xuất diets/foods user đề cập
   - Hoặc tìm các lựa chọn liên quan

2. TẠO BẢNG SO SÁNH:
   Format:
   **[Lựa chọn A]:**
   • Macros: [protein/carb/fat]
   • Ưu điểm: [benefits]
   • Nhược điểm: [drawbacks]
   • Phù hợp cho: [who]
   
   **[Lựa chọn B]:**
   • Macros: [protein/carb/fat]
   • Ưu điểm: [benefits]
   • Nhược điểm: [drawbacks]
   • Phù hợp cho: [who]
   
   **Điểm khác biệt chính:** [key differences]

3. CÂU HỊI TỰ KIỂM TRA:
   Tạo 3-5 câu hỏi giúp user tự đánh giá:
   • Mục tiêu của bạn?
   • Lifestyle như thế nào?
   • Có hạn chế gì không?
   • Thời gian chuẩn bị?

4. Kết thúc: "Bạn trả lời giúp mình để recommend phù hợp nhé!"

QUAN TRỌNG: Dùng RAG knowledge, KHÔNG hard-code."""
        
        # Normal advice
        return """\n\nĐưa ra lời khuyên dinh dưỡng cụ thể, thực tế.
KHÔNG quá lý thuyết.
KHÔNG nói "Dựa trên thông tin"."""
    
    def _handle_general_nutrition_query(self, user_query, chat_history):
        """Handle general nutrition questions using LLM + RAG with comparison support"""
        from config.settings import client, MODEL

        try:
            # Smart RAG - only query when needed (inherit from BaseAgent)
            rag_answer = ''
            rag_sources = []
            
            if self.should_use_rag(user_query, chat_history):
                rag_result = self.rag.query_nutrition(user_query)
                rag_answer = rag_result.get('answer', '')
                rag_sources = rag_result.get('source_docs', [])

            # Build conversation context with RAG context
            rag_context = f"Dựa trên kiến thức từ cơ sở dữ liệu:\n{rag_answer}\n\n" if rag_answer else ""

            messages = [{"role": "system", "content": self.system_prompt}]

            # Add RAG context if available
            if rag_context:
                messages.append({"role": "system", "content": f"Thông tin tham khảo từ cơ sở dữ liệu:\n{rag_context}"})

            # Add chat history (last 5 exchanges)
            if chat_history:
                recent_history = chat_history[-5:] if len(chat_history) > 5 else chat_history
                for user_msg, bot_msg in recent_history:
                    if user_msg:
                        messages.append({"role": "user", "content": user_msg})
                    if bot_msg:
                        messages.append({"role": "assistant", "content": bot_msg})

            # Add current query with context instruction
            context_prompt = self._build_nutrition_context_instruction(user_query, chat_history)
            messages.append({"role": "user", "content": user_query + context_prompt})

            # Get LLM response
            response = client.chat.completions.create(
                model=MODEL,
                messages=messages,
                temperature=0.7,
                max_tokens=500
            )

            llm_response = response.choices[0].message.content

            # Add sources using RAG integration formatter (FIXED!)
            if rag_sources:
                formatted_response = self.rag.format_response_with_sources({
                    'answer': llm_response,
                    'source_docs': rag_sources
                })
                return formatted_response

            return llm_response

        except Exception as e:
            return f"""Xin lỗi, mình gặp lỗi kỹ thuật. Bạn có thể:
1. Thử lại câu hỏi
2. Hỏi cách khác
3. Liên hệ hỗ trợ

Chi tiết lỗi: {str(e)}"""
    
    def should_handoff(self, user_query: str, chat_history: Optional[List] = None) -> bool:
        """
        Override base method - Determine if should hand off to another agent
        
        Specific triggers for nutrition agent:
        - User asks about exercise/workout
        - User mentions symptoms (stomach pain, nausea)
        - User asks about mental health affecting eating
        """
        query_lower = user_query.lower()
        
        # Check each agent's triggers
        for agent, triggers in self.handoff_triggers.items():
            if any(trigger in query_lower for trigger in triggers):
                # Don't handoff if we're in the middle of nutrition consultation
                if chat_history and self._is_mid_consultation(chat_history):
                    return False
                return True
        
        return False
    
    def suggest_next_agent(self, user_query: str) -> Optional[str]:
        """Override base method - Suggest which agent to hand off to based on query"""
        query_lower = user_query.lower()
        
        # Priority order for handoff
        if any(trigger in query_lower for trigger in self.handoff_triggers.get('symptom_agent', [])):
            return 'symptom_agent'
        
        if any(trigger in query_lower for trigger in self.handoff_triggers.get('exercise_agent', [])):
            return 'exercise_agent'
        
        if any(trigger in query_lower for trigger in self.handoff_triggers.get('mental_health_agent', [])):
            return 'mental_health_agent'
        
        if any(trigger in query_lower for trigger in self.handoff_triggers.get('general_health_agent', [])):
            return 'general_health_agent'
        
        return None
    
    def _is_mid_consultation(self, chat_history: List) -> bool:
        """Check if we're in the middle of nutrition consultation"""
        if not chat_history or len(chat_history) < 2:
            return False
        
        # Check last bot response
        last_bot_response = chat_history[-1][1] if len(chat_history[-1]) > 1 else ""
        
        # If we just asked for user data, don't handoff
        if any(phrase in last_bot_response for phrase in [
            "cân nặng", "chiều cao", "tuổi", "giới tính", "mục tiêu"
        ]):
            return True
        
        return False
    
    def _generate_nutrition_summary(self) -> str:
        """Generate summary of nutrition advice for handoff"""
        nutrition_data = self.get_agent_data('nutrition_plan')
        user_profile = self.get_user_profile()
        
        # Natural summary without robotic prefix
        summary_parts = []
        
        if nutrition_data and isinstance(nutrition_data, dict):
            if 'bmi_analysis' in nutrition_data:
                bmi = nutrition_data['bmi_analysis']
                summary_parts.append(f"BMI: {bmi.get('bmi', 'N/A')} ({bmi.get('category', 'N/A')})")
            
            if 'daily_targets' in nutrition_data:
                targets = nutrition_data['daily_targets']
                summary_parts.append(f"Calo: {targets.get('calories', 'N/A')} kcal/ngày")
        
        if user_profile and user_profile.get('goal'):
            summary_parts.append(f"Mục tiêu: {user_profile['goal']}")
        
        return " | ".join(summary_parts)[:100] if summary_parts else ""
    
    def _handle_error(self, error, user_query):
        """Handle errors gracefully"""
        return f"""Xin lỗi, mình gặp chút vấn đề khi tạo tư vấn dinh dưỡng. 😅

Lỗi: {str(error)}

Bạn có thể thử:
1. Cung cấp lại thông tin: tuổi, giới tính, cân nặng, chiều cao
2. Hỏi câu hỏi cụ thể hơn về dinh dưỡng
3. Hoặc mình có thể tư vấn về chủ đề sức khỏe khác

Bạn muốn thử lại không? 💙"""