PamGram/core/template/services.py

import time
from typing import Optional
from urllib.parse import (
    urlencode,
    urljoin,
    urlsplit,
)
from uuid import uuid4

from fastapi import HTTPException
from fastapi.responses import (
    FileResponse,
    HTMLResponse,
)
from fastapi.staticfiles import StaticFiles
from jinja2 import (
    Environment,
    FileSystemLoader,
    Template,
)
from playwright.async_api import ViewportSize

from core.base.aiobrowser import AioBrowser
from core.base.webserver import webapp
from core.bot import bot
from core.template.cache import (
    HtmlToFileIdCache,
    TemplatePreviewCache,
)
from core.template.error import QuerySelectorNotFound
from core.template.models import (
    FileType,
    RenderResult,
)
from utils.const import PROJECT_ROOT
from utils.log import logger


class TemplateService:
    def __init__(
        self,
        browser: AioBrowser,
        html_to_file_id_cache: HtmlToFileIdCache,
        preview_cache: TemplatePreviewCache,
        template_dir: str = "resources",
    ):
        self._browser = browser
        self.template_dir = PROJECT_ROOT / template_dir

        self._jinja2_env = Environment(
            loader=FileSystemLoader(template_dir),
            enable_async=True,
            autoescape=True,
            auto_reload=bot.config.debug,
        )

        self.previewer = TemplatePreviewer(self, preview_cache)

        self.html_to_file_id_cache = html_to_file_id_cache

    def get_template(self, template_name: str) -> Template:
        return self._jinja2_env.get_template(template_name)

    async def render_async(self, template_name: str, template_data: dict) -> str:
        """模板渲染
        :param template_name: 模板文件名
        :param template_data: 模板数据
        """
        start_time = time.time()
        template = self.get_template(template_name)
        html = await template.render_async(**template_data)
        logger.debug(f"{template_name} 模板渲染使用了 {str(time.time() - start_time)}")
        return html

    async def render(
        self,
        template_name: str,
        template_data: dict,
        viewport: Optional[ViewportSize] = None,
        full_page: bool = True,
        evaluate: Optional[str] = None,
        query_selector: Optional[str] = None,
        file_type: FileType = FileType.PHOTO,
        ttl: int = 24 * 60 * 60,
        caption: Optional[str] = None,
        parse_mode: Optional[str] = None,
        filename: Optional[str] = None,
    ) -> RenderResult:
        """模板渲染成图片
        :param template_name: 模板文件名
        :param template_data: 模板数据
        :param viewport: 截图大小
        :param full_page: 是否长截图
        :param evaluate: 页面加载后运行的 js
        :param query_selector: 截图选择器
        :param file_type: 缓存的文件类型
        :param ttl: 缓存时间
        :param caption: 图片描述
        :param parse_mode: 图片描述解析模式
        :param filename: 文件名字
        :return:
        """
        start_time = time.time()
        template = self.get_template(template_name)

        if bot.config.debug:
            preview_url = await self.previewer.get_preview_url(template_name, template_data)
            logger.debug(f"调试模板 URL: {preview_url}")

        html = await template.render_async(**template_data)
        logger.debug(f"{template_name} 模板渲染使用了 {str(time.time() - start_time)}")

        file_id = await self.html_to_file_id_cache.get_data(html, file_type.name)
        if file_id and not bot.config.debug:
            logger.debug(f"{template_name} 命中缓存，返回 file_id {file_id}")
            return RenderResult(
                html=html,
                photo=file_id,
                file_type=file_type,
                cache=self.html_to_file_id_cache,
                ttl=ttl,
                caption=caption,
                parse_mode=parse_mode,
                filename=filename,
            )

        browser = await self._browser.get_browser()
        start_time = time.time()
        page = await browser.new_page(viewport=viewport)
        uri = (PROJECT_ROOT / template.filename).as_uri()
        await page.goto(uri)
        await page.set_content(html, wait_until="networkidle")
        if evaluate:
            await page.evaluate(evaluate)
        clip = None
        if query_selector:
            try:
                card = await page.query_selector(query_selector)
                if not card:
                    raise QuerySelectorNotFound
                clip = await card.bounding_box()
                if not clip:
                    raise QuerySelectorNotFound
            except QuerySelectorNotFound:
                logger.warning(f"未找到 {query_selector} 元素")
        png_data = await page.screenshot(clip=clip, full_page=full_page)
        await page.close()
        logger.debug(f"{template_name} 图片渲染使用了 {str(time.time() - start_time)}")
        return RenderResult(
            html=html,
            photo=png_data,
            file_type=file_type,
            cache=self.html_to_file_id_cache,
            ttl=ttl,
            caption=caption,
            parse_mode=parse_mode,
            filename=filename,
        )


class TemplatePreviewer:
    def __init__(self, template_service: TemplateService, cache: TemplatePreviewCache):
        self.template_service = template_service
        self.cache = cache
        self.register_routes()

    async def get_preview_url(self, template: str, data: dict):
        """获取预览 URL"""
        components = urlsplit(bot.config.webserver.url)
        path = urljoin("/preview/", template)
        query = {}

        # 如果有数据，暂存在 redis 中
        if data:
            key = str(uuid4())
            await self.cache.set_data(key, data)
            query["key"] = key

        return components._replace(path=path, query=urlencode(query)).geturl()

    def register_routes(self):
        """注册预览用到的路由"""

        @webapp.get("/preview/{path:path}")
        async def preview_template(path: str, key: Optional[str] = None):  # pylint: disable=W0612
            # 如果是 /preview/ 开头的静态文件，直接返回内容。比如使用相对链接 ../ 引入的静态资源
            if not path.endswith(".html"):
                full_path = self.template_service.template_dir / path
                if not full_path.is_file():
                    raise HTTPException(status_code=404, detail=f"Template '{path}' not found")
                return FileResponse(full_path)

            # 取回暂存的渲染数据
            data = await self.cache.get_data(key) if key else {}
            if key and data is None:
                raise HTTPException(status_code=404, detail=f"Template data {key} not found")

            # 渲染 jinja2 模板
            html = await self.template_service.render_async(path, data)
            # 将本地 URL file:// 修改为 HTTP url，因为浏览器内不允许加载本地文件
            # file:///project_dir/cache/image.jpg => /cache/image.jpg
            html = html.replace(PROJECT_ROOT.as_uri(), "")
            return HTMLResponse(html)

        # 其他静态资源
        for name in ["cache", "resources"]:
            directory = PROJECT_ROOT / name
            directory.mkdir(exist_ok=True)
            webapp.mount(f"/{name}", StaticFiles(directory=PROJECT_ROOT / name), name=name)