utils/helpers.py

"""
Chat Handler - Uses agent-based architecture with coordination
Clean, modular implementation with specialized agents and memory
"""

from agents import route_to_agent, get_agent, AgentCoordinator
import logging
import os 

os.makedirs("logs", exist_ok=True)

logger = logging.getLogger(__name__)
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('logs/chat_debug.log', mode='a', encoding='utf-8'),
        logging.StreamHandler()
    ]
)

# Setup logging
# logger = logging.getLogger(__name__)
# logging.basicConfig(
#     level=logging.DEBUG,
#     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
#     handlers=[
#         logging.FileHandler('logs/chat_debug.log'),
#         logging.StreamHandler()
#     ]
# )

# Constants
MAX_MESSAGE_LENGTH = 2000
MIN_MESSAGE_LENGTH = 2
SPAM_THRESHOLD_GENTLE = 2
SPAM_THRESHOLD_CONCERNED = 4
SPAM_THRESHOLD_FIRM = 6

# Global coordinator instance (maintains memory across requests)
_coordinator = None

def get_coordinator():
    """Get or create global coordinator instance"""
    global _coordinator
    if _coordinator is None:
        _coordinator = AgentCoordinator()
    return _coordinator


def extract_message_text(message):
    """
    Extract text from message which can be either string or dict (from MultimodalInput)

    Args:
        message: str or dict with 'text' and 'files' keys

    Returns:
        tuple: (text_content, files_list)
    """
    if isinstance(message, dict):
        # MultimodalInput format: {"text": "...", "files": [...]}
        text_content = message.get("text", "").strip()
        files_list = message.get("files", [])
        return text_content, files_list
    elif isinstance(message, str):
        # Regular string message
        return message.strip(), []
    else:
        return "", []


def convert_chatbot_messages_to_list(chat_history):
    """
    Convert ChatbotDataMessage objects to list of [user, bot] pairs

    Args:
        chat_history: List of ChatbotDataMessage objects or list of lists

    Returns:
        list: List of [user_msg, bot_msg] pairs
    """
    if not chat_history:
        return []

    # If already in list format, return as is
    if isinstance(chat_history[0], (list, tuple)):
        logger.debug(f"convert_chatbot_messages_to_list: Already in list format, len={len(chat_history)}")
        return chat_history

    # Convert ChatbotDataMessage objects to list format
    # Messages come as: [user1, bot1, user2, bot2, ...]
    result = []
    i = 0
    logger.debug(f"convert_chatbot_messages_to_list: Converting {len(chat_history)} ChatbotDataMessage objects")

    while i < len(chat_history):
        user_msg = ""
        bot_msg = ""

        # === LẤY USER MESSAGE ===
        if i < len(chat_history):
            item = chat_history[i]
            role = getattr(item, 'role', None) or (item.get('role') if isinstance(item, dict) else None)
            content = getattr(item, 'content', None) or (item.get('content') if isinstance(item, dict) else None)

            if role == 'user':
                user_msg = content
                i += 1
            elif role in ('bot', 'assistant'):
                i += 1
                continue
            else:
                i += 1
                continue

        # === LẤY BOT MESSAGE ===
        if i < len(chat_history):
            item = chat_history[i]
            role = getattr(item, 'role', None) or (item.get('role') if isinstance(item, dict) else None)
            content = getattr(item, 'content', None) or (item.get('content') if isinstance(item, dict) else None)

            if role in ('bot', 'assistant'):
                bot_msg = content
                i += 1

        if user_msg or bot_msg:
            result.append([user_msg, bot_msg])


    logger.debug(f"convert_chatbot_messages_to_list: Result len={len(result)}")
    return result


def convert_list_to_chatbot_messages(chat_history_list):
    """
    Convert list of [user, bot] pairs to ChatbotDataMessage objects with enhanced features

    Args:
        chat_history_list: List of [user_msg, bot_msg] pairs

    Returns:
        list: List of ChatbotDataMessage objects with various features
    """
    from modelscope_studio.components.pro.chatbot import ChatbotDataMessage

    if not chat_history_list:
        return []

    # If already in ChatbotDataMessage format, return as is
    if chat_history_list and hasattr(chat_history_list[0], 'role'):
        logger.debug(f"convert_list_to_chatbot_messages: Already in ChatbotDataMessage format")
        return chat_history_list

    result = []
    logger.debug(f"convert_list_to_chatbot_messages: Converting {len(chat_history_list)} pairs to ChatbotDataMessage")

    for i, (user_msg, bot_msg) in enumerate(chat_history_list):
        # Add user message
        if user_msg:
            result.append(ChatbotDataMessage(
                role="user",
                content=user_msg
            ))

        # Add bot message with enhanced features
        if bot_msg:
            # Determine message features based on content and position
            bot_message_config = {
                "role": "assistant",
                "content": bot_msg
            }

            result.append(ChatbotDataMessage(**bot_message_config))

    logger.debug(f"convert_list_to_chatbot_messages: Result len={len(result)}")
    return result


def create_sample_chatbot_messages():
    """
    Create sample ChatbotDataMessage objects demonstrating various features

    Returns:
        list: List of sample ChatbotDataMessage objects
    """
    from modelscope_studio.components.pro.chatbot import ChatbotDataMessage

    return [
        ChatbotDataMessage(role="user", content="Hello"),
        ChatbotDataMessage(role="assistant", content="World"),
        ChatbotDataMessage(role="assistant",
                           content="Liked message",
                           meta=dict(feedback="like")),
        ChatbotDataMessage(role="assistant",
                           content="Message only has copy button",
                           actions=["copy"]),
        ChatbotDataMessage(
            role="assistant",
            content="Pending message will not show action buttons",
            status="pending"),
        ChatbotDataMessage(
            role="assistant",
            content="Bot 1",
            header="bot1",
            avatar="https://api.dicebear.com/7.x/miniavs/svg?seed=1"),
        ChatbotDataMessage(
            role="assistant",
            content="Bot 2",
            header="bot2",
            avatar="https://api.dicebear.com/7.x/miniavs/svg?seed=2"),
    ]


def chat_logic(message, chat_history, user_id=None):
    """
    Main chat logic using agent routing system

    Args:
        message (str or dict): User's message (string or MultimodalInput dict)
        chat_history (list): List of ChatbotDataMessage objects or [user_msg, bot_msg] pairs
        user_id (str): User ID for data persistence

    Returns:
        tuple: ("", updated_chat_history)
    """

    # ===== INPUT EXTRACTION =====

    # Extract text and files from message (handles both string and dict formats)
    message_text, files_list = extract_message_text(message)

    # Store original message for history
    original_message = message if isinstance(message, str) else message_text

    # Convert ChatbotDataMessage objects to list format if needed
    logger.debug(f"chat_logic input - chat_history type: {type(chat_history)}, len: {len(chat_history) if chat_history else 0}")
    if chat_history and len(chat_history) > 0:
        logger.debug(f"chat_logic input - first item type: {type(chat_history[0])}, has role: {hasattr(chat_history[0], 'role')}")
        if hasattr(chat_history[0], 'content'):
            logger.debug(f"chat_history[0].content: {str(chat_history[0].content)[:50]}")
    chat_history_list = convert_chatbot_messages_to_list(chat_history)
    logger.debug(f"chat_logic after convert - chat_history_list len: {len(chat_history_list)}")
    if chat_history_list:
        logger.debug(f"chat_history_list[0]: {chat_history_list[0]}")

    # ===== INPUT VALIDATION =====

    # Check for empty messages (but allow short acknowledgments like "ờ", "ok", "ừ")
    acknowledgments = ["ờ", "ok", "oke", "ừ", "uhm", "à", "ô", "ồ", "được", "rồi", "vâng", "dạ"]
    if not message_text or (len(message_text) == 0):
        bot_response = "Bạn chưa nhập gì cả. Hãy cho tôi biết bạn cần tư vấn về vấn đề sức khỏe gì nhé! 😊"
        updated_list = chat_history_list + [[original_message, bot_response]]
        updated_chatbot_messages = convert_list_to_chatbot_messages(updated_list)
        return "", updated_chatbot_messages

    # Allow short acknowledgments to pass through
    if message.strip().lower() in acknowledgments:
        # Let the agent handle acknowledgments naturally
        pass  # Continue to agent
    
    # Check for very long messages
    if len(message_text) > 2000:
        bot_response = ("Tin nhắn của bạn quá dài! 😅\n\n"
            "Để tôi có thể tư vấn tốt hơn, hãy chia nhỏ câu hỏi hoặc tóm tắt vấn đề chính của bạn.\n\n"
            "Ví dụ: 'Tôi bị đau đầu 3 ngày, có buồn nôn' thay vì mô tả quá chi tiết.")
        updated_list = chat_history_list + [[original_message, bot_response]]
        updated_chatbot_messages = convert_list_to_chatbot_messages(updated_list)
        return "", updated_chatbot_messages

    # ===== SMART GREETING DETECTION =====
    
    # Detect greeting keywords
    greeting_keywords = [
        "chào", "xin chào", "hello", "hi", "hey", "helo", "hê lô",
        "chao", "alo", "alô", "good morning", "good afternoon", "good evening",
        "buổi sáng", "buổi chiều", "buổi tối", "chào buổi",
        "ê", "ê ơi", "ơi", "ê bot", "ê bạn",  # Vietnamese casual greetings
        "này", "nãy", "nè", "kìa", "ê này"  # More casual Vietnamese
    ]
    
    # Check if message is ONLY a greeting (case-insensitive, strip punctuation)
    message_clean = message.strip().lower().rstrip('!.,?')
    is_pure_greeting = message_clean in greeting_keywords
    
    # Check if it's the first message
    is_first_message = len(chat_history) == 0
    
    # CASE 1: Pure greeting only (e.g., "chào", "hello")
    if is_pure_greeting:
        greeting_response = """Chào bạn! 👋 Mình là trợ lý sức khỏe AI của bạn!

🏥 **Mình có thể giúp gì cho bạn?**

Mình có thể tư vấn về:
• 💊 **Triệu chứng & Sức khỏe** - Phân tích triệu chứng, đề xuất khám bệnh
• 🥗 **Dinh dưỡng** - Lập kế hoạch ăn uống, tính calo, macro
• 💪 **Tập luyện** - Tạo lịch tập gym, hướng dẫn kỹ thuật
• 🧠 **Sức khỏe tâm thần** - Hỗ trợ stress, lo âu, cải thiện giấc ngủ

Bạn đang quan tâm đến vấn đề gì? Hãy chia sẻ với mình nhé! 😊"""
        
        return "", chat_history + [[message, greeting_response]]
    
    # CASE 2: First message with real question (e.g., "đau lưng", "tôi bị đau đầu")
    # Let agent handle it with smart greeting + answer
    
    # ===== SPAM DETECTION =====

    if len(chat_history_list) >= 1:
        all_user_messages = [msg[0] for msg in chat_history_list]
        repeat_count = all_user_messages.count(message_text)
        
        # Level 1: Gentle response (2-3 times)
        if repeat_count == 2:
            bot_response = ("Tôi thấy bạn vừa gửi tin nhắn này lần thứ hai rồi. 😊\n\n"
                "Có phải câu trả lời của tôi chưa giải quyết được vấn đề bạn đang gặp phải không? "
                "Nếu vậy, bạn có thể chia sẻ thêm chi tiết để tôi hiểu rõ hơn không?\n\n"
                "Tôi ở đây để lắng nghe và hỗ trợ bạn. Hãy kể cho tôi nghe thêm nhé! 💙")
            updated_list = chat_history_list + [[original_message, bot_response]]
            updated_chatbot_messages = convert_list_to_chatbot_messages(updated_list)
            return "", updated_chatbot_messages

        # Level 2: Concerned response (4-5 times)
        elif repeat_count >= 4 and repeat_count < 6:
            bot_response = ("Tôi nhận thấy bạn đang lặp lại cùng một câu nhiều lần. Tôi hơi lo lắng - "
                "có phải bạn đang gặp khó khăn trong việc diễn đạt, hay bạn cảm thấy không được lắng nghe?\n\n"
                "Hãy thử cách này nhé:\n"
                "• Nếu bạn đang khó chịu hay đau đớn - hãy mô tả cảm giác đó\n"
                "• Nếu bạn cần thông tin cụ thể - hãy hỏi trực tiếp\n"
                "• Nếu câu trả lời trước không hữu ích - hãy nói cho tôi biết tại sao\n\n"
                "Bạn có muốn bắt đầu lại cuộc trò chuyện không? Tôi sẵn sàng lắng nghe. 🙏")
            updated_list = chat_history_list + [[original_message, bot_response]]
            updated_chatbot_messages = convert_list_to_chatbot_messages(updated_list)
            return "", updated_chatbot_messages

        # Level 3: Firm boundary (6+ times)
        elif repeat_count >= 6:
            bot_response = ("Này, tôi cần nói thẳng với bạn một chút. 😔\n\n"
                "Bạn đã gửi tin nhắn giống nhau " + str(repeat_count) + " lần rồi. "
                "**Nếu bạn thực sự cần giúp đỡ:**\n"
                "Hãy nhấn nút \"Xóa lịch sử\" và bắt đầu lại. Lần này, hãy nói với tôi điều bạn thực sự cần.\n\n"
                "Tôi hy vọng bạn hiểu. Chúc bạn khỏe mạnh! 💚")
            updated_list = chat_history_list + [[original_message, bot_response]]
            updated_chatbot_messages = convert_list_to_chatbot_messages(updated_list)
            return "", updated_chatbot_messages
    
    # ===== AGENT ROUTING =====

    try:
        # Option 1: Use coordinator for memory & multi-agent support (NEW!)
        USE_COORDINATOR = True  # Set to False to use old routing

        if USE_COORDINATOR:
            coordinator = get_coordinator()
            # Pass user_id for data persistence and file analysis
            # Ensure chat_history_list is in correct format for coordinator
            response = coordinator.handle_query(message_text, chat_history_list, user_id=user_id)

            # Convert updated list back to ChatbotDataMessage format
            updated_list = chat_history_list + [[original_message, response]]
            updated_chatbot_messages = convert_list_to_chatbot_messages(updated_list)
            return "", updated_chatbot_messages

        # Option 2: Original routing (fallback)
        # Route to appropriate agent using function calling
        routing_result = route_to_agent(message_text, chat_history_list)

        agent_name = routing_result['agent']
        parameters = routing_result['parameters']

        # Get the specialized agent
        agent = get_agent(agent_name)

        # Let the agent handle the request
        response = agent.handle(parameters, chat_history_list)
        logger.debug(f"Agent {agent_name} response type: {type(response)}")

        # Convert updated list back to ChatbotDataMessage format
        updated_list = chat_history_list + [[original_message, response]]
        updated_chatbot_messages = convert_list_to_chatbot_messages(updated_list)
        return "", updated_chatbot_messages

    except Exception as e:
        # Fallback to general health agent if routing fails
        logger.error(f"Agent routing error: {e}", exc_info=True)

        try:
            from agents.specialized.general_health_agent import GeneralHealthAgent
            agent = GeneralHealthAgent()
            # Ensure chat_history_list is properly formatted
            logger.debug(f"Fallback agent - chat_history_list type: {type(chat_history_list)}, len: {len(chat_history_list)}")
            response = agent.handle({"user_query": message_text}, chat_history_list)
            logger.debug(f"Fallback agent response type: {type(response)}")

            # Convert updated list back to ChatbotDataMessage format
            updated_list = chat_history_list + [[original_message, response]]
            updated_chatbot_messages = convert_list_to_chatbot_messages(updated_list)
            return "", updated_chatbot_messages
        except Exception as e2:
            # Ultimate fallback
            logger.error(f"General health agent error: {e2}", exc_info=True)
            error_response = f"""Xin lỗi, tôi gặp chút vấn đề kỹ thuật. 😅

Lỗi: {str(e2)}

Bạn có thể thử:
1. Hỏi lại câu hỏi
2. Làm mới trang và thử lại
3. Hoặc liên hệ hỗ trợ kỹ thuật

Tôi xin lỗi vì sự bất tiện này! 🙏"""
            # Convert updated list back to ChatbotDataMessage format
            updated_list = chat_history_list + [[original_message, error_response]]
            updated_chatbot_messages = convert_list_to_chatbot_messages(updated_list)
            return "", updated_chatbot_messages