agents/specialized/exercise_agent.py

"""
Exercise Agent - Specialized agent for exercise and fitness advice
"""

from config.settings import client, MODEL
from modules.exercise.exercise import generate_exercise_plan
from health_data import HealthContext
from fitness_tracking import FitnessTracker
from rag.rag_integration import get_rag_integration
from agents.core.base_agent import BaseAgent
from typing import Dict, Any, List, Optional
from datetime import datetime
import re

class ExerciseAgent(BaseAgent):
    def __init__(self, memory=None):
        super().__init__(memory)
        self.health_context = None
        self.fitness_tracker = None
        self.rag = get_rag_integration()
        
        # Configure handoff triggers for exercise agent
        self.handoff_triggers = {
            'nutrition_agent': ['ăn gì', 'thực đơn', 'calo', 'dinh dưỡng', 'giảm cân nhanh', 'tăng cân'],
            'symptom_agent': ['đau', 'chấn thương', 'bị thương', 'sưng', 'viêm'],
            'mental_health_agent': ['stress', 'lo âu', 'không có động lực', 'chán'],
            'general_health_agent': ['khám', 'bác sĩ', 'xét nghiệm']
        }
        self.system_prompt = """Bạn là huấn luyện viên cá nhân chuyên nghiệp, nhiệt huyết và động viên.

💪 CHUYÊN MÔN:
- Tạo kế hoạch tập luyện cá nhân hóa
- Tư vấn bài tập phù hợp với thể trạng, mục tiêu
- Hướng dẫn kỹ thuật tập an toàn
- Tư vấn tập cho người có bệnh nền
- Lịch tập gym, tập tại nhà, cardio, yoga...

🎯 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 30 tuổi, nam, muốn giảm cân, có thể tập 45 phút/ngày" → ĐỪNG HỎI LẠI!
   - Chỉ hỏi thông tin THỰC SỰ còn thiếu
   - Nếu đã đủ thông tin cơ bản → TẠO LỊCH TẬP NGAY!

2. **THÔNG TIN CẦN THIẾT:**
   - Cơ bản: Tuổi, giới tính, mục tiêu, thời gian rảnh
   - Bổ sung: Thể lực, dụng cụ có sẵn, bệnh nền
   - Nếu thiếu → Hỏi ngắn gọn, không hỏi mãi

3. **TẠO LỊCH TẬP:**
   - Lịch tập cụ thể theo ngày
   - Giải thích TẠI SAO tập bài này
   - Hướng dẫn progression (tuần 1, 2, 3...)
   - Lưu ý an toàn, tránh chấn thương

⚠️ AN TOÀN:
- Người có bệnh tim, huyết áp → khuyên gặp bác sĩ trước
- Người có chấn thương → tập nhẹ, tránh vùng bị thương
- Người mới bắt đầu → từ từ, không quá sức

💬 PHONG CÁCH:
- Động viên, khích lệ 💪🔥
- Thực tế, không lý thuyết suông
- Dễ hiểu, dễ làm theo
- Hài hước nhẹ nhà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ĩ, không phải PT!)"""

    def set_health_context(self, health_context: HealthContext):
        """Inject health context and initialize fitness tracker"""
        self.health_context = health_context
        self.fitness_tracker = FitnessTracker(health_context)

    def handle(self, parameters, chat_history=None):
        """
        Handle exercise request

        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 exercise data for next agent
                self.save_agent_data('last_exercise_advice', {
                    'query': user_query,
                    'user_profile': self.get_user_profile(),
                    'timestamp': datetime.now().isoformat()
                })
                
                # Check if nutrition agent shared data with us
                nutrition_data = self.get_other_agent_data('nutrition_agent', 'nutrition_plan')
                context = self._generate_exercise_summary(nutrition_data)
                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,
                'fitness_level': profile.fitness_level,
                'activity_level': profile.activity_level,
                'health_conditions': profile.health_conditions
            }
        # 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 we have enough data - check shared memory first
        profile = self.get_user_profile()
        for field in ['age', 'gender', 'weight', 'height']:
            if not user_data.get(field) and profile.get(field):
                user_data[field] = profile[field]
        
        missing_fields = self._check_missing_data(user_data)

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

        # Generate exercise plan
        try:
            plan = generate_exercise_plan(user_data)

            # Adjust difficulty based on fitness tracker
            if self.fitness_tracker:
                metrics = self.fitness_tracker.calculate_progress_metrics()
                if metrics.get('adherence', 0) > 0.8:
                    plan = self.fitness_tracker.adjust_difficulty(plan, 'increase')
                elif metrics.get('adherence', 0) < 0.5:
                    plan = self.fitness_tracker.adjust_difficulty(plan, 'decrease')

            response = plan

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

            return response
        except Exception as e:
            return self._handle_error(e, user_query)
    
    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,
            'fitness_level': 'beginner',
            'goal': 'health_improvement',
            'available_time': 30,
            '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 fitness level
        if re.search(r'mới bắt đầu|beginner|chưa tập', all_messages.lower()):
            user_data['fitness_level'] = 'beginner'
        elif re.search(r'trung bình|intermediate|tập được', all_messages.lower()):
            user_data['fitness_level'] = 'intermediate'
        elif re.search(r'nâng cao|advanced|tập lâu', all_messages.lower()):
            user_data['fitness_level'] = 'advanced'
        
        # Extract goal
        if re.search(r'giảm cân|weight loss|slim', all_messages.lower()):
            user_data['goal'] = 'weight_loss'
        elif re.search(r'tăng cân|weight gain|bulk', all_messages.lower()):
            user_data['goal'] = 'weight_gain'
        elif re.search(r'tập gym|muscle|cơ bắp|tăng cơ', all_messages.lower()):
            user_data['goal'] = 'muscle_building'
        elif re.search(r'khỏe mạnh|health|sức khỏe', all_messages.lower()):
            user_data['goal'] = 'health_improvement'
        
        # Extract available time
        time_match = re.search(r'(\d+)\s*phút|(\d+)\s*tiếng', all_messages.lower())
        if time_match:
            time_val = int([g for g in time_match.groups() if g][0])
            if 'tiếng' in all_messages.lower():
                time_val *= 60
            user_data['available_time'] = time_val
        
        return user_data
    
    def _check_missing_data(self, user_data):
        """Check what data is missing"""
        required = ['age', 'gender', 'fitness_level', 'goal']
        return [field for field in required if not user_data.get(field)]
    
    def _ask_for_missing_data(self, missing_fields, current_data):
        """Ask for missing data"""
        questions = {
            'age': "bạn bao nhiêu tuổi",
            'gender': "bạn là nam hay nữ",
            'fitness_level': "thể lực hiện tại của bạn thế nào (mới bắt đầu/trung bình/nâng cao)",
            'goal': "mục tiêu của bạn là gì (giảm cân/tăng cơ/khỏe mạnh hơn)"
        }
        
        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ạo lịch tập phù hợp, mình cần biết thêm:**

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

💡 **Ví dụ:** "Tôi 30 tuổi, nam, mới bắt đầu tập, muốn giảm cân, có thể tập 45 phút mỗi ngày"

Sau khi có đủ thông tin, mình sẽ tạo kế hoạch tập luyện 7 ngày chi tiết cho bạn! 🔥"""
    
    def _handle_general_exercise_query(self, user_query, chat_history):
        """Handle general exercise questions using LLM + RAG"""
        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_exercise(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
            messages.append({"role": "user", "content": user_query})

            # 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. Hoặc hỏi mình về chủ đề sức khỏe khác nhé! 💙

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 exercise agent:
        - User asks about nutrition/diet
        - User mentions pain/injury
        - User asks about mental health
        """
        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 exercise planning
                if chat_history and self._is_mid_planning(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"""
        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('nutrition_agent', [])):
            return 'nutrition_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_planning(self, chat_history: List) -> bool:
        """Check if we're in the middle of exercise planning"""
        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 [
            "tuổi", "giới tính", "mục tiêu", "thời gian", "dụng cụ"
        ]):
            return True
        
        return False
    
    def _generate_exercise_summary(self, nutrition_data=None) -> str:
        """Generate summary of exercise advice for handoff"""
        exercise_data = self.get_agent_data('exercise_plan')
        user_profile = self.get_user_profile()
        
        # Natural summary without robotic prefix
        summary_parts = []
        
        if exercise_data and isinstance(exercise_data, dict):
            if 'goal' in exercise_data:
                summary_parts.append(f"Mục tiêu: {exercise_data['goal']}")
            if 'frequency' in exercise_data:
                summary_parts.append(f"Tần suất: {exercise_data['frequency']}")
        
        # Include nutrition data if available (agent-to-agent communication)
        if nutrition_data and isinstance(nutrition_data, dict):
            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('fitness_level'):
            summary_parts.append(f"Thể lực: {user_profile['fitness_level']}")
        
        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 lịch tập. 😅

Lỗi: {str(error)}

Bạn có thể thử:
1. Cung cấp lại thông tin: tuổi, giới tính, thể lực, mục tiêu
2. Hỏi câu hỏi cụ thể hơn về tập luyện
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? 💙"""