"""
消息处理和文件保存工具类
"""
import os
import base64
import uuid
import re
from pathlib import Path
from typing import List, Any, Union


class MessageProcessor:
    """消息处理和文件保存工具类"""
    
    def __init__(self, save_dir: str = None):
        """
        初始化消息处理器
        
        Args:
            save_dir: 文件保存目录，默认为环境变量 ATTACH_SAVE_DIR 或 "uploads"
        """
        self.save_dir = save_dir or os.getenv("ATTACH_SAVE_DIR", "uploads")
        self._saved_files: list[str] = []
    
    def save_and_get_file_url(self, file_data: str, filename: str = None, mime_type: str = None) -> str:
        """
        保存文件并返回绝对路径字符串
        
        Args:
            file_data: base64编码的文件数据
            filename: 文件名
            mime_type: MIME类型
            
        Returns:
            文件绝对路径字符串
        """
        uploads = Path(self.save_dir)
        uploads.mkdir(exist_ok=True)
        
        # 生成安全文件名
        safe_name = filename or f"{uuid.uuid4().hex}"
        if mime_type == "application/pdf" and not safe_name.lower().endswith(".pdf"):
            safe_name += ".pdf"
        
        out_path = uploads / safe_name
        
        try:
            # 解码 base64 数据并保存
            with open(out_path, "wb") as f:
                f.write(base64.b64decode(file_data))
            abs_path = str(out_path.resolve())
            print(f"  -> saved file: {abs_path}")
            self._saved_files.append(abs_path)
            return abs_path
        except Exception as e:
            print(f"  -> save file failed: {e}")
            return f"[附件保存失败: {filename or 'unknown'}]"
    
    def process_content(self, content: Any) -> str:
        """
        处理消息内容，提取文件并转换为文本
        
        Args:
            content: 消息内容，可能是字符串、列表或字典
            
        Returns:
            处理后的文本内容
        """
        if isinstance(content, str):
            return content
        elif isinstance(content, list):
            text_parts = []
            for part in content:
                if isinstance(part, str):
                    text_parts.append(part)
                elif isinstance(part, dict):
                    part_type = part.get("type")
                    if part_type == "text":
                        text_parts.append(part.get("text", ""))
                    elif part_type == "file":
                        # 处理文件类型
                        file_data = part.get("data")
                        filename = part.get("metadata", {}).get("filename")
                        mime_type = part.get("mime_type")
                        
                        if file_data:
                            file_path = self.save_and_get_file_url(file_data, filename, mime_type)
                            text_parts.append(file_path)
                        else:
                            text_parts.append(f"[文件缺失: {filename or 'unknown'}]")
                    else:
                        # 未知类型最小化占位
                        text_parts.append(f"[{part_type or 'unknown'}]")
            return "\n".join([t for t in text_parts if t])
        elif isinstance(content, dict):
            # 单个字典内容
            ctype = content.get("type")
            if ctype == "text":
                return content.get("text", "")
            if ctype == "file":
                file_data = content.get("data")
                filename = content.get("metadata", {}).get("filename")
                mime_type = content.get("mime_type")
                
                if file_data:
                    return self.save_and_get_file_url(file_data, filename, mime_type)
                else:
                    return f"[文件缺失: {filename or 'unknown'}]"
            else:
                return f"[{ctype or 'unknown'}]"
        else:
            return str(content)
    
    def process_messages(self, messages: List[Any]) -> tuple[List[Any], List[str]]:
        """
        处理消息列表，提取文件并转换为文本
        
        Args:
            messages: 消息列表
            
        Returns:
            (处理后的消息列表, 保存的文件路径列表)
        """
        filtered_messages = []
        
        for idx, m in enumerate(messages or []):
            # 兼容 dict 或 LangChain 的消息对象
            if isinstance(m, dict):
                role = m.get("role")
                content = m.get("content")
            else:
                role = getattr(m, "role", None)
                content = getattr(m, "content", None)
            
            print(f"[msg#{idx}] role={role!r}, content_type={type(content).__name__}")

            # 仅当 HumanMessage 且 content 中包含 file 分片时进行处理；否则保持原样
            msg_type = m.get("type") if isinstance(m, dict) else m.__class__.__name__
            is_human = (msg_type == "HumanMessage" or msg_type == "human")
            should_process_file = is_human and isinstance(content, list) and any(isinstance(p, dict) and p.get("type") == "file" for p in content)
            
            # 检查是否包含 waybill_id（非文件处理）
            should_process_waybill = False
            if is_human and isinstance(content, list) and not should_process_file:
                # 检查 content 中是否有 waybill_id 参数
                for part in content:
                    if isinstance(part, dict) and "waybill_id" in part:
                        should_process_waybill = True
                        break

            if should_process_file:
                # 提取文本与文件路径
                text_parts: list[str] = []
                file_paths: list[str] = []
                # 保留的参数：slip_id, token, ver
                preserved_params = {}
                
                for part in content:
                    if isinstance(part, dict):
                        if part.get("type") == "text":
                            text_parts.append(part.get("text", ""))
                            # 从文本部分提取 slip_id, token, ver 参数
                            if "slip_id" in part:
                                preserved_params["slip_id"] = part.get("slip_id")
                            if "token" in part:
                                preserved_params["token"] = part.get("token")
                            if "ver" in part:
                                preserved_params["ver"] = part.get("ver")
                        elif part.get("type") == "file":
                            # 兼容两种文件结构：
                            # 1) {"type":"file", "data":"<base64>", "metadata":{"filename":...}, "mime_type":"application/pdf"}
                            # 2) {"type":"file", "file":{"file_data":"data:application/pdf;base64,<base64>", "filename":"..."}}
                            file_data = part.get("data")
                            filename = part.get("metadata", {}).get("filename")
                            mime_type = part.get("mime_type")

                            if not file_data and isinstance(part.get("file"), dict):
                                f = part.get("file") or {}
                                file_data_url = f.get("file_data")
                                filename = f.get("filename") or filename
                                # 解析 data URL 或原始 base64
                                if isinstance(file_data_url, str):
                                    if file_data_url.startswith("data:") and "," in file_data_url:
                                        try:
                                            header, b64_payload = file_data_url.split(",", 1)
                                            # 格式如 data:application/pdf;base64
                                            if header.startswith("data:") and ";" in header:
                                                mime_type = header[5:].split(";", 1)[0] or mime_type
                                            file_data = b64_payload
                                        except Exception:
                                            file_data = None
                                    else:
                                        # 非 data URL，当作纯 base64
                                        file_data = file_data_url

                            if file_data:
                                saved_path = self.save_and_get_file_url(file_data, filename, mime_type)
                                file_paths.append(saved_path)
                merged_text = "\n".join([t for t in text_parts if t])
                
                # 如果存在 slip_id，将其拼接到文本中
                if "slip_id" in preserved_params:
                    slip_id_value = preserved_params["slip_id"]
                    # 检查文本中是否已有 slip_id：xxx 的格式，如果有则替换，否则追加
                    slip_id_pattern = r'slip_id[：:]\s*\d+'
                    if re.search(slip_id_pattern, merged_text):
                        # 替换原有的 slip_id：xxx
                        merged_text = re.sub(slip_id_pattern, f'slip_id：{slip_id_value}', merged_text)
                    else:
                        # 追加 slip_id：xxx
                        merged_text = f"{merged_text} slip_id：{slip_id_value}"
                
                if file_paths:
                    # 将文件路径拼接为 pdf_path：xxx 的格式，多个文件用逗号分隔
                    pdf_paths_str = "，".join([f"pdf_path：{path}" for path in file_paths])
                    merged_text = f"{merged_text}，{pdf_paths_str}"

                # 将保留的参数添加到content字典中，而不是additional_kwargs
                new_content_dict = {"type": "text", "text": merged_text}
                if preserved_params:
                    new_content_dict.update(preserved_params)
                new_content = [new_content_dict]
                
                # 获取原有的 additional_kwargs（不合并保留的参数）
                if isinstance(m, dict):
                    existing_kwargs = m.get("additional_kwargs", {})
                else:
                    existing_kwargs = getattr(m, "additional_kwargs", {})
                
                # 保持原有的 additional_kwargs，不添加保留的参数
                merged_kwargs = existing_kwargs

                if isinstance(m, dict):
                    filtered_msg = {**m, "content": new_content, "additional_kwargs": merged_kwargs}
                else:
                    try:
                        filtered_msg = m.__class__(
                            content=new_content,
                            additional_kwargs=merged_kwargs,
                            response_metadata=getattr(m, "response_metadata", {}),
                            id=getattr(m, "id", None),
                        )
                    except Exception:
                        # 兜底为等价字典并保留 id
                        filtered_msg = {
                            "type": m.__class__.__name__,
                            "role": (role or "user"),
                            "id": getattr(m, "id", None),
                            "additional_kwargs": merged_kwargs,
                            "response_metadata": getattr(m, "response_metadata", {}),
                            "content": new_content,
                        }
                filtered_messages.append(filtered_msg)
                try:
                    print(f"  -> processed to: {len(merged_text)} chars")
                    if preserved_params:
                        print(f"  -> preserved params: {preserved_params}")
                except Exception:
                    print("  -> processed")
                    if preserved_params:
                        print(f"  -> preserved params: {preserved_params}")
            elif should_process_waybill:
                # 处理包含 waybill_id 的消息
                new_content_list = []
                for part in content:
                    if isinstance(part, dict):
                        if part.get("type") == "text" and "waybill_id" in part:
                            # 提取文本和 waybill_id
                            text = part.get("text", "")
                            waybill_id = part.get("waybill_id")
                            
                            # 将 waybill_id 拼接到 text 中
                            waybill_pattern = r'waybill_id[：:]\s*\d+'
                            if re.search(waybill_pattern, text):
                                # 如果文本中已有 waybill_id：xxx，则替换
                                text = re.sub(waybill_pattern, f'waybill_id：{waybill_id}', text)
                            else:
                                # 否则追加 waybill_id：xxx
                                text = f"{text} waybill_id：{waybill_id}"
                            
                            # 创建新的 content 字典，保留 waybill_id 字段
                            new_part = {
                                "type": "text",
                                "text": text,
                                "waybill_id": waybill_id
                            }
                            new_content_list.append(new_part)
                        else:
                            # 其他部分保持不变
                            new_content_list.append(part)
                    else:
                        new_content_list.append(part)
                
                # 更新消息
                if isinstance(m, dict):
                    filtered_msg = {**m, "content": new_content_list}
                else:
                    try:
                        filtered_msg = m.__class__(
                            content=new_content_list,
                            additional_kwargs=getattr(m, "additional_kwargs", {}),
                            response_metadata=getattr(m, "response_metadata", {}),
                            id=getattr(m, "id", None),
                        )
                    except Exception:
                        filtered_msg = {
                            "type": m.__class__.__name__,
                            "role": (role or "user"),
                            "id": getattr(m, "id", None),
                            "additional_kwargs": getattr(m, "additional_kwargs", {}),
                            "response_metadata": getattr(m, "response_metadata", {}),
                            "content": new_content_list,
                        }
                filtered_messages.append(filtered_msg)
                print(f"  -> processed waybill_id message")
            else:
                # 不处理，其它消息保持不变
                filtered_messages.append(m)
                print("  -> processed (no change)")
        
        return filtered_messages, list(self._saved_files)