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
        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

        loop = asyncio.get_event_loop()
        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:]}")

        # --- SỬA LỖI LOGIC CHỐNG TRÙNG LẶP TẠI ĐÂY ---
        # Kiểm tra xem timestamp của sự kiện webhook này đã tồn tại trong lịch sử chưa
        for row in history:
            # Chuyển đổi an toàn sang string để so sánh
            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 # Bỏ qua hoàn toàn để tránh xử lý lại

        # --- LUỒNG XỬ LÝ GỐC CỦA BẠN ĐƯỢC GIỮ NGUYÊN ---
        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}")

        # Thêm timestamp mới nếu chưa có (logic này có thể không cần thiết nữa nhưng giữ lại để không thay đổi luồng)
        conv['timestamp'] = self.flatten_timestamp(conv['timestamp'])
        if timestamp not in conv['timestamp']:
            conv['timestamp'].append(timestamp)
        
        # Lần gọi thứ 2 để cập nhật thêm thông tin ban đầu (nếu có)
        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:
            await 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 = None
        cau_hoi = None
        if isinstance(llm_analysis, dict):
            keywords = [self.normalize_vehicle_keyword(llm_analysis.get('phuong_tien', ''))]
            muc_dich = llm_analysis.get('muc_dich')
            tu_khoa = llm_analysis.get('tu_khoa')
            cau_hoi = llm_analysis.get('cau_hoi')
        elif isinstance(llm_analysis, list) and len(llm_analysis) > 0:
            keywords = [self.normalize_vehicle_keyword(llm_analysis[0].get('phuong_tien', ''))]
            muc_dich = llm_analysis[0].get('muc_dich')
            tu_khoa = llm_analysis[0].get('tu_khoa')
            cau_hoi = llm_analysis[0].get('cau_hoi')
        else:
            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()
        
        logger.info(f"[DEBUG] Phương tiện: {keywords} - Hành vi: {tu_khoa} - Mục đích: {muc_dich} - Câu hỏi: {cau_hoi}")

        # Hợp nhất dữ liệu đã phân tích vào `conv`
        conv['originalcommand'] = command
        conv['originalcontent'] = remaining_text
        conv['originalvehicle'] = ','.join(keywords)
        conv['originalaction'] = tu_khoa
        conv['originalpurpose'] = muc_dich
        conv['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
        if not command:
            if muc_dich_to_use == "hỏi về mức phạt":
                response = await self.handle_muc_phat(conv, conversation_context, page_token, sender_id)
            elif muc_dich_to_use == "hỏi về quy tắc giao thông":
                response = await self.handle_quy_tac(conv, conversation_context, page_token, sender_id)
            elif muc_dich_to_use == "hỏi về báo hiệu đường bộ":
                response = await self.handle_bao_hieu(conv, conversation_context, page_token, sender_id)
            elif muc_dich_to_use == "hỏi về quy trình xử lý vi phạm giao thông":
                response = await self.handle_quy_trinh(conv, conversation_context, page_token, sender_id)
            elif muc_dich_to_use == "thông tin cá nhân của AI":
                response = await self.handle_ca_nhan(conv, conversation_context, page_token, sender_id)
            else:
                response = await self.handle_khac(conv, conversation_context, page_token, sender_id)
        else:
            if command == "xong":
                post_url = await self.create_facebook_post(page_token, conv['recipient_id'], [conv])
                if post_url:
                    response = f"Bài viết đã được tạo thành công! Bạn có thể xem tại: {post_url}"
                else:
                    response = "Đã xảy ra lỗi khi tạo bài viết. Vui lòng thử lại sau."
                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

        await 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}")
        await 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):
        sorted_history = sorted(history, key=lambda row: self.get_latest_timestamp(row.get('timestamp', 0)))

        total_chars = 0
        MAX_CONTEXT_CHARS = 20_000
        conversation_context = []
        for row in reversed(sorted_history):
            temp_blocks = []
            
            # --- SỬA LỖI THỨ TỰ LỊCH SỬ TẠI ĐÂY ---
            # Đảm bảo lượt nói của user luôn được thêm vào trước.
            if row.get('originaltext'):
                temp_blocks.append({"role": "user", "content": row['originaltext']})
            # Lượt nói của trợ lý ảo được thêm vào sau.
            if row.get('systemresponse'):
                temp_blocks.append({"role": "assistant", "content": row['systemresponse']})
            
            temp_total = sum(len(block['content']) for block in temp_blocks)
            if total_chars + temp_total > MAX_CONTEXT_CHARS: continue
            
            # Thêm cặp hỏi-đáp vào đầu danh sách context, duy trì thứ tự thời gian
            conversation_context = temp_blocks + conversation_context
            total_chars += temp_total
        return conversation_context

    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."
        await self.facebook.send_message(message=get_random_message(FOUND_REGULATIONS_MESSAGES))
        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}"
            impounding = match.get('impounding')
            if impounding:
                full_result_text += f"\n- Có thể tạm giữ phương tiện."
        
        # --- PROMPT MỚI ĐƯỢC CẬP NHẬT TẠI ĐÂ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:"
        )
        
        await 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()}")
        
        # Fallback response
        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}")
        # In a real scenario, you would use the Facebook Graph API to create a post.
        # This is a mock implementation.
        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)})")
            
            match_count = get_settings().match_count
            matches = 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}")
            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:
        """Hàm chung để xử lý các câu hỏi kiến thức chung."""
        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_general_question(conversation_context, conv['originaltext'], "quy tắc giao thông")
        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_general_question(conversation_context, conv['originaltext'], "báo hiệu đường bộ")
        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_general_question(conversation_context, conv['originaltext'], "quy trình xử lý vi phạm giao thông")
        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'
            # --- SỬA LỖI CÚ PHÁP TẠI ĐÂY ---
            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!"