Agents/note_agent/note_formatter.py

"""
笔记格式化模块

将多模态内容（图片、文本等）转换为 Markdown 笔记格式。
支持可扩展的处理器架构，方便后续添加音频、视频、PDF 等格式。
"""

import base64
import hashlib
import logging
from abc import ABC, abstractmethod
from datetime import datetime
from pathlib import Path
from typing import Optional

logger = logging.getLogger("adk.agent")


# ============================================================
# 内容处理器基类
# ============================================================
class BaseContentHandler(ABC):
    """内容处理器基类，所有格式处理器继承此类。"""

    @abstractmethod
    def can_handle(self, mime_type: str) -> bool:
        """判断是否能处理该 MIME 类型"""
        ...

    @abstractmethod
    def process(self, data: bytes, mime_type: str, filename: str = "") -> str:
        """
        处理内容并返回 Markdown 片段。

        Args:
            data: 原始字节数据
            mime_type: MIME 类型
            filename: 文件名（可选）

        Returns:
            Markdown 格式的字符串片段
        """
        ...


# ============================================================
# 图片处理器
# ============================================================
class ImageHandler(BaseContentHandler):
    """处理图片类型，保存到本地并生成 Markdown 图片引用。"""

    SUPPORTED_TYPES = {"image/png", "image/jpeg", "image/jpg", "image/gif", "image/webp", "image/svg+xml"}

    def can_handle(self, mime_type: str) -> bool:
        return mime_type.lower() in self.SUPPORTED_TYPES

    def process(self, data: bytes, mime_type: str, filename: str = "") -> str:
        # 确保输出目录存在
        NOTES_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

        # 生成文件名
        if not filename:
            ext = self._mime_to_ext(mime_type)
            filename = f"image_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{hashlib.md5(data).hexdigest()[:8]}.{ext}"

        filepath = NOTES_OUTPUT_DIR / filename
        filepath.write_bytes(data)

        logger.info("图片已保存: %s", filepath)
        return f"![{filename}](notes/{filename})"

    @staticmethod
    def _mime_to_ext(mime_type: str) -> str:
        return {
            "image/png": "png",
            "image/jpeg": "jpg",
            "image/jpg": "jpg",
            "image/gif": "gif",
            "image/webp": "webp",
            "image/svg+xml": "svg",
        }.get(mime_type.lower(), "bin")


# ============================================================
# 文本处理器
# ============================================================
class TextHandler(BaseContentHandler):
    """处理纯文本内容。"""

    SUPPORTED_TYPES = {"text/plain", "text/markdown", "text/csv"}

    def can_handle(self, mime_type: str) -> bool:
        return mime_type.lower() in self.SUPPORTED_TYPES

    def process(self, data: bytes, mime_type: str, filename: str = "") -> str:
        text = data.decode("utf-8", errors="replace")
        if mime_type.lower() == "text/markdown":
            return text  # 已经是 Markdown，直接返回
        return text  # 纯文本直接返回


# ============================================================
# 占位处理器（用于尚未支持的格式）
# ============================================================
class PlaceholderHandler(BaseContentHandler):
    """占位处理器，为尚未支持的格式生成占位标记。"""

    def can_handle(self, mime_type: str) -> bool:
        return True  # 兜底，处理所有未匹配的类型

    def process(self, data: bytes, mime_type: str, filename: str = "") -> str:
        name = filename or "unsupported_file"
        return f"<!-- [{name}] 格式 {mime_type} 暂不支持，待后续扩展 -->"


# ============================================================
# 笔记格式化器
# ============================================================
class NoteFormatter:
    """
    笔记格式化器，管理所有内容处理器。

    用法：
        formatter = NoteFormatter()
        md_content = formatter.format_to_markdown(
            title="会议笔记",
            text_content="讨论了项目进度",
            files=[(image_bytes, "image/png", "screenshot.png")],
        )
    """

    def __init__(self):
        self._handlers: list[BaseContentHandler] = [
            ImageHandler(),
            TextHandler(),
            PlaceholderHandler(),  # 必须放最后，作为兜底
        ]

    def add_handler(self, handler: BaseContentHandler):
        """注册新的内容处理器（插入到占位处理器之前）"""
        # 移除末尾的 PlaceholderHandler
        if self._handlers and isinstance(self._handlers[-1], PlaceholderHandler):
            self._handlers.pop()
        self._handlers.append(handler)
        # 重新添加 PlaceholderHandler 到末尾
        self._handlers.append(PlaceholderHandler())

    def _find_handler(self, mime_type: str) -> BaseContentHandler:
        for handler in self._handlers:
            if handler.can_handle(mime_type):
                return handler
        return self._handlers[-1]  # PlaceholderHandler

    def format_to_markdown(
        self,
        title: str = "",
        text_content: str = "",
        files: Optional[list[tuple[bytes, str, str]]] = None,
    ) -> str:
        """
        将多模态内容格式化为 Markdown 笔记。

        Args:
            title: 笔记标题
            text_content: 文本内容
            files: 文件列表，每项为 (data, mime_type, filename)

        Returns:
            完整的 Markdown 笔记字符串
        """
        lines: list[str] = []
        now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        # 标题
        if title:
            lines.append(f"# {title}")
            lines.append("")
        else:
            lines.append(f"# 笔记 - {now}")
            lines.append("")

        # 元信息
        lines.append(f"> 创建时间：{now}")
        lines.append("")

        # 文本内容
        if text_content:
            lines.append("## 内容")
            lines.append("")
            lines.append(text_content)
            lines.append("")

        # 附件
        if files:
            lines.append("## 附件")
            lines.append("")
            for data, mime_type, filename in files:
                handler = self._find_handler(mime_type)
                try:
                    md_fragment = handler.process(data, mime_type, filename)
                    lines.append(md_fragment)
                    lines.append("")
                except Exception as e:
                    logger.error("处理文件失败: %s, 错误: %s", filename, e)
                    lines.append(f"<!-- 文件 {filename} 处理失败: {e} -->")
                    lines.append("")

        return "\n".join(lines)