app/message_processor.py

from typing import Dict, Any, List, Optional
import asyncio
import traceback
import json
from loguru import logger
from .constants import SUMMARY_STATUS_MESSAGES, PROCESSING_STATUS_MESSAGES, FOUND_REGULATIONS_MESSAGES, BATCH_STATUS_MESSAGES
from .utils import get_random_message
from .facebook import FacebookClient
from app.config import get_settings
import re

class MessageProcessor:
    def __init__(self, channel, sender_id):
        self.channel = channel
        self.sender_id = sender_id
        # FacebookClient riêng cho từng conversation
        self.facebook = FacebookClient(
            app_secret=get_settings().facebook_app_secret,
            page_id=channel.page_id,
            page_token=channel.get_page_token(),
            sender_id=sender_id
        )

    async def process_message(self, message_data: Dict[str, Any]):
        if not message_data or not isinstance(message_data, dict):
            logger.error(f"[ERROR] Invalid message_data: {message_data}")
            return
        required_fields = ["sender_id", "page_id", "text", "timestamp"]
        for field in required_fields:
            if field not in message_data:
                logger.error(f"[ERROR] Missing field {field} in message_data: {message_data}")
                return
        
        loop = asyncio.get_event_loop()
        sender_id = message_data["sender_id"]
        page_id = message_data["page_id"]
        message_text = message_data["text"]
        timestamp = message_data["timestamp"]
        attachments = message_data.get('attachments', [])
        logger.bind(user_id=sender_id, page_id=page_id, message=message_text).info("Processing message")

        if not message_text and not attachments:
            logger.info(f"[DEBUG] Không có message_text và attachments, không xử lý...")
            return

        sheets_client = self.channel.get_sheets_client()
        history = await loop.run_in_executor(
            None, lambda: sheets_client.get_conversation_history(sender_id, page_id)
        )
        logger.info(f"[DEBUG] history: ... {history[-3:]}")

        for row in history:
            sheet_timestamps = [str(ts) for ts in row.get('timestamp', [])]
            if str(timestamp) in sheet_timestamps:
                logger.warning(f"Webhook lặp lại cho sự kiện đã tồn tại (timestamp: {timestamp}). Bỏ qua.")
                return

        log_kwargs = {
            'conversation_id': None, 'recipient_id': sender_id, 'page_id': page_id,
            'originaltext': message_text, 'originalcommand': '', 'originalcontent': '',
            'originalattachments': attachments, 'originalvehicle': '', 'originalaction': '',
            'originalpurpose': '', 'originalquestion': '', 'systemresponse': '',
            'timestamp': [timestamp], 'isdone': False
        }

        logger.info(f"[DEBUG] Message cơ bản: {log_kwargs}")
        conv = await loop.run_in_executor(None, lambda: sheets_client.log_conversation(**log_kwargs))
        if not conv:
            logger.error("Không thể tạo conversation mới!")
            return
        logger.info(f"[DEBUG] Message history sau lần ghi đầu: {conv}")
                
        conv['timestamp'] = self.flatten_timestamp(conv['timestamp'])
        if timestamp not in conv['timestamp']:
            conv['timestamp'].append(timestamp)
        
        conv_after_update1 = await loop.run_in_executor(None, lambda: sheets_client.log_conversation(**conv))
        if conv_after_update1:
            conv = conv_after_update1

        page_token = self.channel.get_page_token()
        if not page_token:
            logger.error(f"No access token found for page {message_data['page_id']}")
            return
        
        try:
            asyncio.create_task(self.facebook.send_message(message=get_random_message(PROCESSING_STATUS_MESSAGES)))
        except Exception as e:
            if "expired" in str(e).lower():
                logger.warning("[FACEBOOK] Token expired, invalidate and refresh")
                self.channel.invalidate_page_token()
                page_token = self.channel.get_page_token(force_refresh=True)
                self.facebook.page_token = page_token
            else:
                raise

        from app.utils import extract_command, extract_keywords
        from app.constants import VEHICLE_KEYWORDS
        command, remaining_text = extract_command(message_text)
        
        llm_analysis = await self.channel.llm.analyze(message_text, self.get_llm_history(history))
        logger.info(f"[LLM][RAW] Kết quả trả về từ analyze: {llm_analysis}")
        
        muc_dich = None
        tu_khoa_list = [] # Sửa: đổi tên thành tu_khoa_list và khởi tạo là list rỗng
        cau_hoi = None
        
        # Sửa: Đơn giản hóa logic, vì LLM giờ luôn trả về 1 dict
        analysis_data = None
        if isinstance(llm_analysis, list) and llm_analysis:
            analysis_data = llm_analysis[0]
        elif isinstance(llm_analysis, dict):
            analysis_data = llm_analysis

        if analysis_data:
            # Lấy phương tiện và chuẩn hóa
            phuong_tien = self.normalize_vehicle_keyword(analysis_data.get('phuong_tien', ''))
            keywords = [phuong_tien] if phuong_tien else []
            
            muc_dich = analysis_data.get('muc_dich')
            
            # Lấy danh sách từ khóa, đảm bảo nó là list
            raw_tu_khoa = analysis_data.get('tu_khoa', [])
            if isinstance(raw_tu_khoa, list):
                tu_khoa_list = raw_tu_khoa
            elif isinstance(raw_tu_khoa, str) and raw_tu_khoa:
                tu_khoa_list = [raw_tu_khoa] # Chuyển string thành list 1 phần tử
                
            cau_hoi = analysis_data.get('cau_hoi')
        else:
            # Fallback logic cũ nếu LLM không phân tích được
            keywords = extract_keywords(message_text, VEHICLE_KEYWORDS)
            cau_hoi = message_text
            for kw in keywords: cau_hoi = cau_hoi.replace(kw, "")
            cau_hoi = cau_hoi.strip()
        
        # Sửa: Log danh sách từ khóa
        logger.info(f"[DEBUG] Phương tiện: {keywords} - Từ khóa pháp lý: {tu_khoa_list} - Mục đích: {muc_dich} - Câu hỏi: {cau_hoi}")

        conv.update({
            'originalcommand': command, 'originalcontent': remaining_text, 'originalvehicle': ','.join(keywords),
            'originalaction': ' '.join(tu_khoa_list), 'originalpurpose': muc_dich, 'originalquestion': cau_hoi or ""
        })

        muc_dich_to_use = muc_dich or conv.get('originalpurpose')
        logger.info(f"[DEBUG] Định hướng mục đích xử lý: {muc_dich_to_use}")
        conversation_context = self.get_llm_history(history)

        response = None
        handlers = {
            "hỏi về mức phạt": self.handle_muc_phat,
            "hỏi về quy tắc giao thông": self.handle_quy_tac,
            "hỏi về báo hiệu đường bộ": self.handle_bao_hieu,
            "hỏi về quy trình xử lý vi phạm giao thông": self.handle_quy_trinh,
            "thông tin cá nhân của AI": self.handle_ca_nhan
        }

        if not command:
            handler = handlers.get(muc_dich_to_use, self.handle_khac)
            response = await handler(conv, conversation_context, page_token, sender_id)
        else:
            if command == "xong":
                post_url = await self.create_facebook_post(page_token, conv['recipient_id'], [conv])
                response = f"Bài viết đã được tạo thành công! Bạn có thể xem tại: {post_url}" if post_url else "Đã xảy ra lỗi khi tạo bài viết."
                conv['isdone'] = True
            else:
                response = "Vui lòng cung cấp thêm thông tin và gõ lệnh \\xong khi hoàn tất."
                conv['isdone'] = False

        asyncio.create_task(self.facebook.send_message(message=response))
        
        conv['systemresponse'] = response
        
        logger.info(f"Chuẩn bị ghi/cập nhật dữ liệu cuối cùng vào sheet: {conv}")
        
        loop.run_in_executor(None, lambda: sheets_client.log_conversation(**conv))
        
        return
    
    def get_latest_timestamp(self,ts_value):
        if isinstance(ts_value, (int, float)): return int(ts_value)
        if isinstance(ts_value, str):
            try: return int(json.loads(ts_value))
            except: 
                try: return int(ts_value)
                except: return 0
        if isinstance(ts_value, list):
            if not ts_value: return 0
            return max([self.get_latest_timestamp(item) for item in ts_value]) if ts_value else 0
        return 0
    
    def get_llm_history(self, history: List[Dict[str, Any]]) -> str:
        """
        Định dạng lịch sử hội thoại thành một chuỗi văn bản duy nhất,
        bao gồm cả các từ khóa đã sử dụng để cung cấp ngữ cảnh cho LLM.
        """
        sorted_history = sorted(history, key=lambda row: self.get_latest_timestamp(row.get('timestamp', 0)))
        
        # Lấy 5 lượt hội thoại gần nhất để tránh context quá dài
        recent_history = sorted_history[-5:]
        
        context_lines = []
        for row in recent_history:
            user_text = row.get('originaltext', '').strip()
            assistant_text = row.get('systemresponse', '').strip()
            keywords_used = row.get('originalaction', '').strip()

            if user_text:
                context_lines.append(f"Người dùng: {user_text} (từ khóa đã dùng: {keywords_used})")
            
            if assistant_text:
                context_lines.append(f"Trợ lý: {assistant_text}")
                
        return "\n".join(context_lines)

    def flatten_timestamp(self, ts):
        flat = []
        if not isinstance(ts, list): ts = [ts]
        for t in ts:
            if isinstance(t, list): flat.extend(self.flatten_timestamp(t))
            else: flat.append(t)
        return flat

    def normalize_vehicle_keyword(self, keyword: str) -> str:
        from app.constants import VEHICLE_KEYWORDS
        import difflib
        if not keyword: return ""
        matches = difflib.get_close_matches(keyword.lower(), [k.lower() for k in VEHICLE_KEYWORDS], n=1, cutoff=0.6)
        if matches:
            for k in VEHICLE_KEYWORDS:
                if k.lower() == matches[0]: return k
        return keyword
        
    async def format_search_results(self, conversation_context: str, question: str, matches: List[Dict[str, Any]], page_token: str, sender_id: str) -> str:
        if not matches:
            return "Không tìm thấy kết quả phù hợp."

        asyncio.create_task(self.facebook.send_message(message=get_random_message(FOUND_REGULATIONS_MESSAGES)))

        #TODO: thời gian rerank kéo dài hơn 30s. Tạm thời bỏ qua bước reranking cho đến khi tìm ra phương án optimize 
        # try:
        #     reranked = await self.channel.reranker.rerank(question, matches, top_k=10)
        #     if reranked: matches = reranked
        # except Exception as e:
        #     logger.error(f"[RERANK] Lỗi khi rerank: {e}")
        
        full_result_text = ""
        def arr_to_str(arr, sep=", "):
            if not arr: return ""
            return sep.join([str(x) for x in arr if x not in (None, "")]) if isinstance(arr, list) else str(arr)

        for i, match in enumerate(matches, 1):
            full_result_text += f"\n- Nguồn: {(match.get('structure') or '').strip()}:\n"
            fullContent = (match.get('fullcontent') or '').strip()
            full_result_text += f"{fullContent}"
            hpbsnoidung = arr_to_str(match.get('hpbsnoidung'), sep="; ")
            if hpbsnoidung: full_result_text += f"\n- Hình phạt bổ sung: {hpbsnoidung}"
            bpkpnoidung = arr_to_str(match.get('bpkpnoidung'), sep="; ")
            if bpkpnoidung: full_result_text += f"\n- Biện pháp khắc phục: {bpkpnoidung}"
            if match.get('impounding'): full_result_text += f"\n- Tạm giữ phương tiện: 07 ngày"
        
        prompt = (
            "Bạn là một trợ lý pháp lý AI chuyên nghiệp. Nhiệm vụ của bạn là tổng hợp thông tin từ hai nguồn: **Lịch sử trò chuyện** và **Các đoạn luật liên quan** để đưa ra một câu trả lời duy nhất, liền mạch và tự nhiên cho người dùng.\n\n"
            "**QUY TẮC BẮT BUỘC:**\n"
            "1.  **Hành văn tự nhiên:** Trả lời thẳng vào câu hỏi. **Không** bắt đầu bằng các cụm từ như 'Dựa trên thông tin được cung cấp', 'Theo các đoạn luật', v.v.\n"
            "2.  **Nguồn trích dẫn:** Khi cần trích dẫn, chỉ nêu nguồn từ văn bản luật (ví dụ: 'theo Khoản 1, Điều 5...'). **Tuyệt đối không** trích dẫn nguồn là 'từ lịch sử trò chuyện'.\n"
            "3.  **Tổng hợp thông tin:** Phải kết hợp thông tin từ cả hai nguồn một cách mượt mà. Ví dụ, nếu lịch sử trò chuyện đã có mức phạt cho xe máy, và câu hỏi hiện tại là về xe máy điện, hãy sử dụng thông tin từ văn bản luật để xác định xe máy điện thuộc nhóm xe nào, sau đó áp dụng mức phạt đã biết từ lịch sử.\n"
            "4.  **Ngắn gọn, chính xác:** Luôn trả lời ngắn gọn, rõ ràng và chỉ dựa vào thông tin được cung cấp.\n\n"
            f"### Lịch sử trò chuyện:\n{conversation_context}\n\n"
            f"### Các đoạn luật liên quan:\n{full_result_text}\n\n"
            f"### Câu hỏi của người dùng:\n{question}\n\n"
            "### Trả lời:"
        )
        
        asyncio.create_task(self.facebook.send_message(message=f"{get_random_message(SUMMARY_STATUS_MESSAGES)}"))
        
        try:
            answer = await self.channel.llm.generate_text(prompt)
            if answer and answer.strip():
                logger.info(f"LLM trả về câu trả lời: \n\tanswer: {answer}")
                return answer.strip()
            else:
                logger.error(f"LLM không trả về câu trả lời phù hợp: \n\tanswer: {answer}")
        except Exception as e:
            logger.error(f"LLM không sẵn sàng: {e}\n{traceback.format_exc()}")
        
        return "Dựa trên thông tin bạn cung cấp, tôi đã tìm thấy một số quy định liên quan. Tuy nhiên, tôi đang gặp chút khó khăn trong việc tóm tắt. Bạn vui lòng tham khảo nội dung chi tiết trong các văn bản luật nhé."

    async def create_facebook_post(self, page_token: str, sender_id: str, history: List[Dict[str, Any]]) -> str:
        logger.info(f"[MOCK] Creating Facebook post for sender_id={sender_id} with history={history}")
        return "https://facebook.com/mock_post_url" 

    async def handle_muc_phat(self, conv, conversation_context, page_token, sender_id):
        vehicle = conv.get('originalvehicle', '')
        action = conv.get('originalaction', '')
        question = conv.get('originalquestion', '')

        
        if not action and not question:
             return "Để tra cứu mức phạt, bạn vui lòng cung cấp hành vi vi phạm nhé."

        search_query = action or question
        logger.info(f"[DEBUG] tạo embedding cho: '{search_query}'")
        try:
            embedding = await self.channel.embedder.create_embedding(search_query)
            logger.info(f"[DEBUG] embedding: {embedding[:5]} ... (total {len(embedding)})")
            
            loop = asyncio.get_event_loop()
            match_count = get_settings().match_count
            

            matches = await loop.run_in_executor(
                None,
                lambda: self.channel.supabase.match_documents(
                    embedding,
                    match_count=match_count,
                    user_question=search_query
                )
            )

            logger.info(f"[DEBUG] matches: {matches[:2]}...{matches[-2:]}")
            if matches:
                response = await self.format_search_results(conversation_context, question or action, matches, page_token, sender_id)
            else:
                response = "Xin lỗi, tôi không tìm thấy thông tin phù hợp với hành vi bạn mô tả."
        except Exception as e:
            logger.error(f"Lỗi khi tra cứu mức phạt: {e}\n{traceback.format_exc()}")
            response = "Đã có lỗi xảy ra trong quá trình tra cứu. Vui lòng thử lại sau."
            
        conv['isdone'] = True
        return response

    async def _handle_general_question(self, conversation_context: str, message_text: str, topic: str) -> str:
        prompt = (
            "Bạn là một trợ lý AI am hiểu về luật giao thông Việt Nam. "
            "Dựa vào lịch sử trò chuyện và kiến thức của bạn, hãy trả lời câu hỏi của người dùng một cách rõ ràng, ngắn gọn và chính xác.\n"
            f"Chủ đề câu hỏi là về: {topic}\n"
            f"### Lịch sử:\n{conversation_context}\n"
            f"### Câu hỏi của người dùng:\n{message_text}\n"
            "### Trả lời:"
        )
        try:
            answer = await self.channel.llm.generate_text(prompt)
            if answer and answer.strip():
                return answer.strip()
            return f"Tôi chưa có thông tin về câu hỏi của bạn liên quan đến {topic}."
        except Exception as e:
            logger.error(f"Lỗi khi xử lý chủ đề {topic}: {e}")
            return f"Xin lỗi, tôi đang gặp sự cố khi xử lý câu hỏi về {topic}. Vui lòng thử lại sau."

    async def handle_khac(self, conv, conversation_context, page_token, sender_id):
        conv['isdone'] = True
        return await self._handle_general_question(conversation_context, conv['originaltext'], "một vấn đề khác")

    async def handle_quy_tac(self, conv, conversation_context, page_token, sender_id):
        conv['isdone'] = True
        return await self.handle_muc_phat(conv, conversation_context, page_token, sender_id)

    async def handle_bao_hieu(self, conv, conversation_context, page_token, sender_id):
        conv['isdone'] = True
        return await self.handle_muc_phat(conv, conversation_context, page_token, sender_id)

    async def handle_quy_trinh(self, conv, conversation_context, page_token, sender_id):
        conv['isdone'] = True
        return await self.handle_muc_phat(conv, conversation_context, page_token, sender_id)

    async def handle_ca_nhan(self, conv, conversation_context, page_token, sender_id):
        prompt = (
            "Biết rằng bạn đã có lịch sử trao đổi như sau:"
            f"Lịch sử:\n{conversation_context}\n\n"
            'Với các thông tin sau: "Bạn có tên là WeThoong AI, là trợ lý giao thông thông minh. Bạn được anh Viet Cat tạo ra và facebook cá nhân của anh ấy là https://facebook.com/vietcat". '
            'Không được trả lời bạn là AI của Google, OpenAI, hay bất kỳ hãng nào khác. '
            'Hãy trả lời thông minh, hài hước, ngắn gọn cho câu hỏi sau:\n'
            f'Câu hỏi:\n"{conv["originaltext"]}"'
        )
        try:
            answer = await self.channel.llm.generate_text(prompt)
            conv['isdone'] = True
            return answer.strip() if answer and answer.strip() else "Chào bạn, mình là WeThoong AI đây!"
        except Exception as e:
            logger.error(f"Lỗi khi xử lý câu hỏi cá nhân: {e}")
            return "Chào bạn, mình là WeThoong AI, trợ lý giao thông thông minh của bạn!"