Blame for src/email_parser.py

"""
邮件解析模块 - 解析 .eml 文件，提取关键字段
支持 UTF-8、GBK、GB2312、ISO-8859-1 等编码
"""
import io
import re
import logging
import base64
from dataclasses import dataclass, field
from datetime import datetime
from email import policy
from email.header import decode_header as email_decode_header
from email.message import Message
from pathlib import Path
from typing import Optional, List

import html2text

logger = logging.getLogger(__name__)

# 常见需要解码的邮件编码列表（按优先级）
COMMON_ENCODINGS = ["utf-8", "gbk", "gb2312", "gb18030", "iso-8859-1", "windows-1252"]


@dataclass
class AttachmentInfo:
    """附件信息"""
    filename: str
    content_type: str
    size: int  # 字节数


@dataclass
class ParsedEmail:
    """解析后的邮件数据结构"""
    filename: str
    sender: str
    sender_domain: str
    recipient: str = ""
    recipient_domain: str = ""
    subject: str = ""
    date: Optional[datetime] = None
    date_str: str = ""
    body_text: str = ""
    body_html: str = ""
    body_preview: str = ""  # 前200字符，用于展示
    attachments: List[AttachmentInfo] = field(default_factory=list)
    raw_path: str = ""
    parse_error: str = ""

    @property
    def is_valid(self) -> bool:
        return not self.parse_error


def _decode_header_value(header_value: Optional[str]) -> str:
    """
    解码邮件头字段（如 Subject、From、To）。
    处理 RFC 2047 编码字，如 =?UTF-8?B?xxxxx?=
    """
    if not header_value:
        return ""
    decoded_parts = []
    for part_bytes, encoding in email_decode_header(header_value):
        if isinstance(part_bytes, bytes):
            decoded_parts.append(_decode_bytes(part_bytes, encoding))
        else:
            decoded_parts.append(str(part_bytes))
    return "".join(decoded_parts).strip()


def _decode_bytes(data: bytes, suggested_encoding: Optional[str] = None) -> str:
    """
    尝试用多种编码解码字节数据。
    优先使用建议的编码，然后依次尝试常见编码。
    """
    if not data:
        return ""

    # 如果指定了编码，先试它
    if suggested_encoding:
        try:
            return data.decode(suggested_encoding, errors="strict")
        except (UnicodeDecodeError, LookupError):
            pass

    # 依次尝试常见编码
    for enc in COMMON_ENCODINGS:
        try:
            return data.decode(enc, errors="strict")
        except (UnicodeDecodeError, LookupError):
            continue

    # 全部失败，用 ignore 模式强制解码
    return data.decode("utf-8", errors="ignore")


def _extract_sender_domain(email_addr: str) -> str:
    """从邮箱地址提取主域名"""
    match = re.search(r'@([\w.\-]+)', email_addr)
    if match:
        domain = match.group(1).lower()
        parts = domain.split(".")
        if len(parts) >= 2:
            return parts[-2]
        return parts[0]
    return ""


def _extract_email_address(header_value: str) -> str:
    """
    从邮件头（如 'Alice <alice@example.com>'）中提取纯邮箱地址。
    """
    if not header_value:
        return ""
    # 优先匹配尖括号中的地址
    match = re.search(r'<([^>]+)>', header_value)
    if match:
        return match.group(1).lower().strip()
    # 否则直接匹配邮箱格式
    match = re.search(r'[\w.+\-]+@[\w.\-]+', header_value)
    if match:
        return match.group(0).lower()
    return header_value.strip()


def _html_to_text(html_content: str) -> str:
    """将 HTML 正文转为纯文本"""
    if not html_content:
        return ""
    try:
        converter = html2text.HTML2Text()
        converter.ignore_links = True
        converter.ignore_images = True
        converter.ignore_emphasis = False
        converter.body_width = 0  # 不换行
        text = converter.handle(html_content)
        # 清理多余空行
        lines = [line.rstrip() for line in text.splitlines()]
        cleaned = "\n".join(line for line in lines if line.strip())
        return cleaned
    except Exception as e:
        logger.warning(f"HTML 转文本失败: {e}")
        # 降级：用正则去掉 HTML 标签
        return re.sub(r'<[^>]+>', ' ', html_content)


def _clean_body(text: str) -> str:
    """清理邮件正文，去除多余空白"""
    if not text:
        return ""
    text = text.replace('\r\n', '\n').replace('\r', '\n')
    text = re.sub(r'\n{3,}', '\n\n', text)
    return text.strip()


def _decode_part_payload(part) -> str:
    """
    解码邮件 MIME 部分的 payload。
    处理 base64、quoted-printable 等传输编码，
    以及各类字符编码（UTF-8、GBK 等）。
    """
    payload = part.get_payload(decode=True)
    if not payload:
        return ""

    # 获取声明的字符编码
    charset = part.get_content_charset() or None
    return _decode_bytes(payload, charset)


def _parse_msg_common(msg: Message, filename: str) -> ParsedEmail:
    """
    从已解析的 email.Message 对象中提取字段（供文件和字节流共用）。
    """
    # 解码各项头字段
    raw_sender = _decode_header_value(msg.get("From", ""))
    raw_recipient = _decode_header_value(msg.get("To", ""))
    subject = _decode_header_value(msg.get("Subject", ""))

    sender = _extract_email_address(raw_sender)
    recipient = _extract_email_address(raw_recipient)
    sender_domain = _extract_sender_domain(sender)
    recipient_domain = _extract_sender_domain(recipient)

    # 日期
    date_obj = None
    date_str = ""
    date_header = msg.get("Date", "")
    if date_header:
        try:
            from email.utils import parsedate_to_datetime
            date_obj = parsedate_to_datetime(date_header)
            if date_obj:
                date_str = date_obj.strftime("%Y-%m-%d")
        except Exception as e:
            logger.warning(f"日期解析失败 ({filename}): {e}")
            # 尝试直接用原始字符串
            date_str = date_header.strip()

    # 正文和附件
    body_text = ""
    body_html = ""
    attachments = []

    if msg.is_multipart():
        for part in msg.walk():
            content_type = part.get_content_type()
            content_disposition = part.get("Content-Disposition", "")

            # 处理附件
            if "attachment" in content_disposition or \
               (content_type not in ("text/plain", "text/html") and part.get_filename()):
                filename_att = _decode_header_value(part.get_filename())
                if not filename_att:
                    filename_att = f"attachment_{len(attachments) + 1}"
                payload = part.get_payload(decode=True)
                attachments.append(AttachmentInfo(
                    filename=filename_att,
                    content_type=content_type,
                    size=len(payload) if payload else 0,
                ))
                continue

            # 处理正文
            if content_type == "text/plain" and not body_text:
                body_text = _decode_part_payload(part)
            elif content_type == "text/html" and not body_html:
                body_html = _decode_part_payload(part)
    else:
        # 单部分邮件
        content_type = msg.get_content_type()
        if content_type == "text/html":
            body_html = _decode_part_payload(msg)
        else:
            body_text = _decode_part_payload(msg)

    # 如果只有 HTML 正文，转换为纯文本
    if not body_text and body_html:
        body_text = _clean_body(_html_to_text(body_html))
    elif body_text:
        body_text = _clean_body(body_text)

    # 如果两者都为空
    if not body_text and not body_html:
        logger.warning(f"邮件正文为空: {filename}")
        body_text = "[邮件正文为空]"

    body_preview = body_text[:200].replace('\n', ' ') + ("..." if len(body_text) > 200 else "")

    return ParsedEmail(
        filename=filename,
        sender=sender,
        sender_domain=sender_domain,
        recipient=recipient,
        recipient_domain=recipient_domain,
        subject=subject,
        date=date_obj,
        date_str=date_str,
        body_text=body_text,
        body_html=body_html,
        body_preview=body_preview,
        attachments=attachments,
        raw_path="",
    )


def parse_eml_file(file_path: str) -> ParsedEmail:
    """
    解析单个 .eml 文件

    Args:
        file_path: .eml 文件路径

    Returns:
        ParsedEmail 对象
    """
    filename = Path(file_path).name

    try:
        from email import message_from_binary_file
        with open(file_path, "rb") as f:
            msg = message_from_binary_file(f, policy=policy.default)

        result = _parse_msg_common(msg, filename)
        result.raw_path = file_path
        return result

    except Exception as e:
        logger.error(f"解析 .eml 文件失败 ({filename}): {e}", exc_info=True)
        return ParsedEmail(
            filename=filename,
            sender="",
            sender_domain="",
            subject="",
            body_text="",
            body_preview="",
            raw_path=file_path,
            parse_error=str(e),
        )


def parse_eml_bytes(file_bytes: bytes, filename: str) -> ParsedEmail:
    """
    从字节流解析 .eml 文件（Gradio 上传的场景）

    Args:
        file_bytes: 文件字节内容
        filename: 原始文件名

    Returns:
        ParsedEmail 对象
    """
    try:
        from email import message_from_bytes
        msg = message_from_bytes(file_bytes, policy=policy.default)
        result = _parse_msg_common(msg, filename)
        return result

    except Exception as e:
        logger.error(f"解析 .eml 字节流失败 ({filename}): {e}", exc_info=True)
        return ParsedEmail(
            filename=filename,
            sender="",
            sender_domain="",
            subject="",
            body_text="",
            body_preview="",
            parse_error=str(e),
        )


def batch_parse_eml_files(file_paths: list[str]):
    """批量解析多个 .eml 文件"""
    success = []
    failed = []
    for path in file_paths:
        parsed = parse_eml_file(path)
        if parsed.is_valid:
            success.append(parsed)
        else:
            failed.append(parsed)
    return {"success": success, "failed": failed}