Blame for src/email_parser.py

"""
邮件解析模块 - 解析 .msg 文件，提取关键字段
"""
import io
import re
import logging
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Optional

import extract_msg
import html2text

logger = logging.getLogger(__name__)


@dataclass
class ParsedEmail:
    """解析后的邮件数据结构"""
    filename: str
    sender: str
    sender_domain: str
    date: Optional[datetime]
    date_str: str
    subject: str
    body_text: str
    body_preview: str  # 前200字符，用于展示
    raw_path: str = ""
    parse_error: str = ""

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


@dataclass
class EmailParseResult:
    """批量解析结果"""
    success: list = field(default_factory=list)
    failed: list = field(default_factory=list)

    @property
    def total(self) -> int:
        return len(self.success) + len(self.failed)


def _extract_sender_domain(sender: str) -> str:
    """从发件人邮箱地址提取域名"""
    match = re.search(r'@([\w.\-]+)', sender)
    if match:
        domain = match.group(1).lower()
        # 去掉顶级域名，取主域名（如 homelight.com -> homelight）
        parts = domain.split(".")
        if len(parts) >= 2:
            return parts[-2]
        return parts[0]
    return ""


def _html_to_text(html_content: str) -> str:
    """将 HTML 正文转为纯文本"""
    converter = html2text.HTML2Text()
    converter.ignore_links = True
    converter.ignore_images = True
    converter.ignore_emphasis = False
    converter.body_width = 0  # 不换行
    try:
        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')
    # 去除连续3个以上空行
    text = re.sub(r'\n{3,}', '\n\n', text)
    return text.strip()


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

    Args:
        file_path: .msg 文件路径（或字节数据文件路径）

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

    try:
        msg = extract_msg.Message(file_path)

        # 发件人
        sender_raw = msg.sender or ""
        # 尝试提取纯邮箱地址
        email_match = re.search(r'[\w.+\-]+@[\w.\-]+', sender_raw)
        sender = email_match.group(0).lower() if email_match else sender_raw.strip()
        sender_domain = _extract_sender_domain(sender)

        # 日期
        date_obj = None
        date_str = ""
        try:
            if msg.date:
                if isinstance(msg.date, datetime):
                    date_obj = msg.date
                else:
                    from dateutil import parser as dateutil_parser
                    date_obj = dateutil_parser.parse(str(msg.date))
                date_str = date_obj.strftime("%Y-%m-%d")
        except Exception as e:
            logger.warning(f"日期解析失败 ({filename}): {e}")
            date_str = str(msg.date) if msg.date else ""

        # 主题
        subject = (msg.subject or "").strip()

        # 正文：优先使用纯文本，否则转换 HTML
        body_text = ""
        if msg.body:
            body_text = _clean_body(msg.body)

        if not body_text and msg.htmlBody:
            html_content = msg.htmlBody
            if isinstance(html_content, bytes):
                html_content = html_content.decode("utf-8", errors="ignore")
            body_text = _clean_body(_html_to_text(html_content))

        # 若两者都为空，记录警告
        if not body_text:
            logger.warning(f"邮件正文为空: {filename}")
            body_text = "[邮件正文为空]"

        # 预览（前200字符）
        body_preview = body_text[:200].replace('\n', ' ') + ("..." if len(body_text) > 200 else "")

        msg.close()

        return ParsedEmail(
            filename=filename,
            sender=sender,
            sender_domain=sender_domain,
            date=date_obj,
            date_str=date_str,
            subject=subject,
            body_text=body_text,
            body_preview=body_preview,
            raw_path=file_path,
        )

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


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

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

    Returns:
        ParsedEmail 对象
    """
    import tempfile
    import os

    # 写入临时文件（extract_msg 需要文件路径）
    with tempfile.NamedTemporaryFile(suffix=".msg", delete=False) as tmp:
        tmp.write(file_bytes)
        tmp_path = tmp.name

    try:
        result = parse_msg_file(tmp_path)
        result.filename = filename  # 恢复原始文件名
        return result
    finally:
        try:
            os.unlink(tmp_path)
        except Exception:
            pass


def batch_parse_msg_files(file_paths: list[str]) -> EmailParseResult:
    """
    批量解析多个 .msg 文件

    Args:
        file_paths: .msg 文件路径列表

    Returns:
        EmailParseResult 对象（含成功/失败列表）
    """
    result = EmailParseResult()
    for path in file_paths:
        parsed = parse_msg_file(path)
        if parsed.is_valid:
            result.success.append(parsed)
        else:
            result.failed.append(parsed)
    return result