NeoBot/core/managers/image_manager.py

"""
图片生成管理器模块

负责管理图片生成相关的逻辑，支持多种渲染引擎（目前支持 Playwright）。
"""
import os
import base64
import tempfile
from typing import Dict, Any, Optional
from jinja2 import Template

from .browser_manager import browser_manager
from ..utils.logger import logger
from ..utils.singleton import Singleton

class ImageManager(Singleton):
    """
    图片生成管理器（单例）
    """

    def __init__(self):
        """
        初始化图片生成管理器
        """
        # 检查是否已经初始化
        if hasattr(self, 'template_dir'):
            return
            
        # 模板目录
        self.template_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "templates")
        # 临时文件目录 - 使用系统临时目录
        self.temp_dir = os.path.join(tempfile.gettempdir(), "neobot_images")
        os.makedirs(self.temp_dir, exist_ok=True)
        # 模板缓存
        self._template_cache: Dict[str, Template] = {}

    async def render_template(self, template_name: str, data: Dict[str, Any], output_name: str = "output.png", quality: int = 80, image_type: str = "png") -> Optional[str]:
        """
        使用 Playwright 渲染 Jinja2 模板并保存为图片文件

        Args:
            template_name (str): 模板文件名 (例如 "help.html")
            data (Dict[str, Any]): 传递给模板的数据字典
            output_name (str, optional): 输出文件名. Defaults to "output.png".
            quality (int, optional): JPEG 质量 (0-100). 仅在 image_type 为 jpeg 时有效. Defaults to 80.
            image_type (str, optional): 图片类型 ('png' or 'jpeg'). Defaults to "png".

        Returns:
            Optional[str]: 生成图片的绝对路径，如果失败则返回 None
        """
        template_path = os.path.join(self.template_dir, template_name)
        if not os.path.exists(template_path):
            logger.error(f"模板文件未找到: {template_path}")
            return None

        try:
            # 1. 渲染 HTML (使用缓存)
            if template_name in self._template_cache:
                template = self._template_cache[template_name]
            else:
                with open(template_path, "r", encoding="utf-8") as f:
                    template_str = f.read()
                template = Template(template_str)
                self._template_cache[template_name] = template
            
            html_content = template.render(**data)

            # 2. 使用浏览器截图
            # 改为从池中获取页面
            page = await browser_manager.get_page()
            if not page:
                logger.error("无法获取浏览器页面")
                return None

            try:
                width = 1920
                height = 1080
                await page.set_viewport_size({"width": width, "height": height})
                
                # 加载内容
                await page.set_content(html_content)
                await page.wait_for_selector("body")
                
                
                screenshot_args = {
                    'full_page': True, 
                    'type': image_type,
                    'omit_background': False,
                    'scale': 'css'
                }
                if image_type == 'jpeg':
                    screenshot_args['quality'] = quality
                    
                screenshot_bytes = await page.screenshot(**screenshot_args)  # type: ignore
                
            finally:
                # 归还页面到池中，而不是直接关闭
                await browser_manager.release_page(page)

            # 3. 保存文件
            output_path = os.path.join(self.temp_dir, output_name)
            with open(output_path, "wb") as f:
                f.write(screenshot_bytes)
                
            logger.info(f"图片已生成: {output_path} ({len(screenshot_bytes)/1024:.2f} KB)")
            return os.path.abspath(output_path)

        except Exception as e:
            logger.exception(f"渲染模板 {template_name} 失败: {e}")
            return None

    async def render_template_to_base64(self, template_name: str, data: Dict[str, Any], output_name: str = "output.png", quality: int = 80, image_type: str = "png") -> Optional[str]:
        """
        渲染模板并返回 Base64 编码的图片字符串
        """
        file_path = await self.render_template(template_name, data, output_name, quality, image_type)
        if not file_path:
            return None
            
        try:
            with open(file_path, "rb") as f:
                content = f.read()
            
            mime_type = "image/jpeg" if image_type == "jpeg" else "image/png"
            base64_str = base64.b64encode(content).decode("utf-8")
            
            # 记录摘要日志，避免刷屏
            log_message = f"Base64 图片已生成 (MIME: {mime_type}, Size: {len(base64_str)/1024:.2f} KB, Preview: {base64_str[:30]}...{base64_str[-30:]})"
            logger.debug(log_message)
            
            return f"data:{mime_type};base64," + base64_str
        except Exception as e:
            logger.error(f"读取图片文件失败: {e}")
            return None

# 全局图片管理器实例
image_manager = ImageManager()