zhenxun_bot/zhenxun/services/renderer/engine.py

import asyncio
from collections import OrderedDict
from collections.abc import Awaitable, Callable
import contextlib
import hashlib
import inspect
import json
from pathlib import Path
import time
from typing import Any, ClassVar, cast

import nonebot_plugin_htmlrender.browser as htmlrender_browser
import psutil

from zhenxun.configs.config import Config
from zhenxun.services.log import logger

from .types import BaseScreenshotEngine


async def _await_if_needed(value: Any) -> Any:
    if inspect.isawaitable(value):
        return await cast(Awaitable[Any], value)
    return value


async def _get_browser_instance() -> Any:
    for attr_name in ("get_browser", "get_new_browser"):
        browser_getter = getattr(htmlrender_browser, attr_name, None)
        if callable(browser_getter):
            return await _await_if_needed(browser_getter())
    raise RuntimeError("nonebot_plugin_htmlrender.browser 未提供可用浏览器获取函数。")


async def _shutdown_browser_instance() -> None:
    for attr_name in (
        "shutdown_htmlrender",
        "shutdown_browser",
        "close_browser",
        "close_htmlrender",
    ):
        shutdown_func = getattr(htmlrender_browser, attr_name, None)
        if callable(shutdown_func):
            try:
                await _await_if_needed(shutdown_func())
            finally:
                with contextlib.suppress(Exception):
                    setattr(htmlrender_browser, "_browser", None)
                with contextlib.suppress(Exception):
                    setattr(htmlrender_browser, "_playwright", None)
            return

    browser_obj = getattr(htmlrender_browser, "_browser", None)
    close_func = getattr(browser_obj, "close", None) if browser_obj else None
    if callable(close_func):
        with contextlib.suppress(Exception):
            await _await_if_needed(close_func())

    playwright_obj = getattr(htmlrender_browser, "_playwright", None)
    stop_func = getattr(playwright_obj, "stop", None) if playwright_obj else None
    if callable(stop_func):
        with contextlib.suppress(Exception):
            await _await_if_needed(stop_func())

    with contextlib.suppress(Exception):
        setattr(htmlrender_browser, "_browser", None)
    with contextlib.suppress(Exception):
        setattr(htmlrender_browser, "_playwright", None)

    if callable(close_func) or callable(stop_func):
        return

    logger.debug(
        "未找到 htmlrender 浏览器关闭函数，跳过 shutdown。",
        "PlaywrightEngine",
    )


def _patch_playwright_env_check_once() -> None:
    if getattr(htmlrender_browser, "_zhenxun_check_once_patched", False):
        return

    original_check: Callable[..., Awaitable[Any]] | None = None
    check_attr_name = ""
    for attr_name in ("check_playwright_env", "check_browser_env"):
        candidate = getattr(htmlrender_browser, attr_name, None)
        if callable(candidate):
            original_check = cast(Callable[..., Awaitable[Any]], candidate)
            check_attr_name = attr_name
            break

    if original_check is None:
        logger.debug(
            "未找到 htmlrender 环境检查函数，跳过 check_once 补丁。",
            "PlaywrightEngine",
        )
        setattr(htmlrender_browser, "_zhenxun_check_once_patched", True)
        return

    check_func = original_check
    state: dict[str, Any] = {"checked": False, "result": None}
    check_lock: asyncio.Lock | None = None

    def _is_browser_usable(browser_obj: Any) -> bool:
        if browser_obj is None:
            return False
        is_connected = getattr(browser_obj, "is_connected", None)
        if callable(is_connected):
            with contextlib.suppress(Exception):
                return bool(is_connected())
        # 无法判断连接状态时，保守认为可用
        return True

    def _get_current_browser_candidate() -> Any:
        current = state["result"]
        if _is_browser_usable(current):
            return current
        fallback = getattr(htmlrender_browser, "_browser", None)
        if _is_browser_usable(fallback):
            return fallback
        return None

    async def _check_once(**kwargs: Any) -> Any:
        nonlocal check_lock
        if state["checked"]:
            cached_browser = _get_current_browser_candidate()
            if cached_browser is not None:
                return cached_browser
            state["checked"] = False
            state["result"] = None

        if check_lock is None:
            check_lock = asyncio.Lock()
        async with check_lock:
            if state["checked"]:
                cached_browser = _get_current_browser_candidate()
                if cached_browser is not None:
                    return cached_browser
                state["checked"] = False
                state["result"] = None

            result = await check_func(**kwargs)
            state["checked"] = True
            state["result"] = result

            browser = _get_current_browser_candidate()
            if browser is not None:
                return browser
            return result

    setattr(htmlrender_browser, check_attr_name, _check_once)
    setattr(htmlrender_browser, "_zhenxun_check_once_patched", True)


class PlaywrightEngine(BaseScreenshotEngine):
    """使用 nonebot-plugin-htmlrender 实现的截图引擎。"""

    _MAX_CONCURRENT_RENDER = 2
    _CONTEXT_POOL_SIZE = 2
    _PREWARM_CONTEXT_COUNT = 1
    _SET_CONTENT_WAIT_UNTIL = "domcontentloaded"
    _READY_STATE_TIMEOUT_MS = 2_000
    _IMAGE_READY_TIMEOUT_MS = 1_800
    _FONT_READY_TIMEOUT_MS = 1_200
    _FULL_PAGE_VIEWPORT_MAX_HEIGHT = 4_096
    _FULL_PAGE_VIEWPORT_MAX_WIDTH = 4_096
    _CLIP_PADDING_DEFAULT = 0
    _DISABLE_ANIMATIONS_STYLE = """
        *, *::before, *::after {
            animation: none !important;
            transition: none !important;
            caret-color: transparent !important;
            scroll-behavior: auto !important;
        }
    """
    _RECENT_RESULT_TTL_SECONDS = 1.5
    _RECENT_RESULT_MAX_ITEMS = 64
    _RSS_RECYCLE_MIN_THRESHOLD_BYTES = 700 * 1024 * 1024
    _RSS_RECYCLE_MAX_THRESHOLD_BYTES = 1200 * 1024 * 1024
    _RSS_RECYCLE_HEADROOM_BYTES = 224 * 1024 * 1024
    _RECYCLE_COOLDOWN_SECONDS = 300
    _RECYCLE_CHECK_EVERY = 8
    _IDLE_CHECK_INTERVAL_SECONDS = 15
    _IDLE_RECYCLE_SECONDS = 180
    _POOL_UNSAFE_OPTION_KEYS: ClassVar[set[str]] = {
        "device_scale_factor",
        "color_scheme",
        "extra_http_headers",
        "forced_colors",
        "geolocation",
        "has_touch",
        "http_credentials",
        "ignore_https_errors",
        "is_mobile",
        "java_script_enabled",
        "locale",
        "permissions",
        "proxy",
        "record_har_content",
        "record_har_mode",
        "record_har_omit_content",
        "record_har_path",
        "record_video_dir",
        "record_video_size",
        "reduced_motion",
        "screen",
        "service_workers",
        "storage_state",
        "timezone_id",
        "user_agent",
    }

    def __init__(self):
        _patch_playwright_env_check_once()
        self._render_semaphore = asyncio.Semaphore(self._MAX_CONCURRENT_RENDER)
        self._debug_console_log = bool(Config.get_config("UI", "DEBUG_MODE", False))
        self._state_lock = asyncio.Lock()
        self._recycle_lock = asyncio.Lock()
        self._active_renders = 0
        self._render_count = 0
        self._recycle_pending = False
        self._last_recycle_at = 0.0
        self._last_render_finished_at = time.monotonic()
        self._rss_baseline_bytes: int | None = None
        self._recent_results: OrderedDict[str, tuple[float, bytes]] = OrderedDict()
        self._inflight_tasks: dict[str, asyncio.Task[bytes]] = {}
        self._context_pool: asyncio.LifoQueue[Any] = asyncio.LifoQueue()
        self._all_contexts: set[Any] = set()
        self._idle_recycle_task: asyncio.Task[None] | None = None
        self._closing = False
        self._process = psutil.Process()

    @staticmethod
    def _normalize_base_url(path: Path) -> str:
        base_url = path.absolute().as_uri()
        if not base_url.endswith("/"):
            base_url += "/"
        return base_url

    @staticmethod
    def _build_render_key(
        html: str, template_path: str, render_options: dict[str, Any]
    ) -> str:
        options_json = json.dumps(render_options, sort_keys=True, default=str)
        hasher = hashlib.sha256()
        hasher.update(template_path.encode("utf-8"))
        hasher.update(b"\0")
        hasher.update(options_json.encode("utf-8"))
        hasher.update(b"\0")
        hasher.update(html.encode("utf-8", errors="ignore"))
        return hasher.hexdigest()

    def _cleanup_recent_results_nolock(self, now: float) -> None:
        while self._recent_results:
            expire_at, _ = next(iter(self._recent_results.values()))
            if expire_at > now:
                break
            self._recent_results.popitem(last=False)
        while len(self._recent_results) > self._RECENT_RESULT_MAX_ITEMS:
            self._recent_results.popitem(last=False)

    def _get_recent_result_nolock(self, key: str, now: float) -> bytes | None:
        entry = self._recent_results.get(key)
        if not entry:
            return None
        expire_at, result = entry
        if expire_at <= now:
            self._recent_results.pop(key, None)
            return None
        self._recent_results.move_to_end(key)
        return result

    def _get_total_rss(self) -> int | None:
        try:
            total_rss = self._process.memory_info().rss
            for child in self._process.children(recursive=True):
                with contextlib.suppress(Exception):
                    total_rss += child.memory_info().rss
            return total_rss
        except Exception:
            return None

    def _update_rss_baseline_nolock(self, current_rss: int) -> None:
        if self._rss_baseline_bytes is None or current_rss < self._rss_baseline_bytes:
            self._rss_baseline_bytes = current_rss
            return

        threshold = self._rss_baseline_bytes + self._RSS_RECYCLE_HEADROOM_BYTES * 2
        if current_rss >= threshold:
            self._rss_baseline_bytes = int(
                self._rss_baseline_bytes * 0.9 + current_rss * 0.1
            )

    def _get_dynamic_threshold_nolock(self, current_rss: int) -> int:
        self._update_rss_baseline_nolock(current_rss)
        baseline = self._rss_baseline_bytes or current_rss
        dynamic = baseline + self._RSS_RECYCLE_HEADROOM_BYTES
        dynamic = max(dynamic, self._RSS_RECYCLE_MIN_THRESHOLD_BYTES)
        dynamic = min(dynamic, self._RSS_RECYCLE_MAX_THRESHOLD_BYTES)
        return dynamic

    def _mark_recycle_if_needed_nolock(self, now: float) -> None:
        if self._render_count % self._RECYCLE_CHECK_EVERY != 0:
            return
        if now - self._last_recycle_at < self._RECYCLE_COOLDOWN_SECONDS:
            return
        current_rss = self._get_total_rss()
        if current_rss is None:
            return
        threshold = self._get_dynamic_threshold_nolock(current_rss)
        if current_rss >= threshold:
            self._recycle_pending = True

    async def initialize(self) -> None:
        async with self._state_lock:
            if self._idle_recycle_task and not self._idle_recycle_task.done():
                return
            self._closing = False
            self._last_render_finished_at = time.monotonic()
            if current_rss := self._get_total_rss():
                self._rss_baseline_bytes = current_rss
            self._idle_recycle_task = asyncio.create_task(self._idle_recycle_loop())
        await self._prewarm_browser_and_pool()

    async def close(self) -> None:
        idle_task: asyncio.Task[None] | None = None
        async with self._state_lock:
            self._closing = True
            idle_task = self._idle_recycle_task
            self._idle_recycle_task = None
            for task in self._inflight_tasks.values():
                task.cancel()
            self._inflight_tasks.clear()
            self._recent_results.clear()
            self._recycle_pending = False

        if idle_task:
            idle_task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await idle_task

        await self._dispose_context_pool()
        await _shutdown_browser_instance()

    async def _on_render_begin(self) -> None:
        async with self._state_lock:
            self._active_renders += 1

    async def _on_render_end(self) -> None:
        should_recycle = False
        async with self._state_lock:
            self._active_renders = max(0, self._active_renders - 1)
            self._render_count += 1
            now = time.monotonic()
            self._last_render_finished_at = now
            self._mark_recycle_if_needed_nolock(now)
            if self._recycle_pending and self._active_renders == 0:
                self._recycle_pending = False
                self._last_recycle_at = now
                should_recycle = True
        if should_recycle:
            await self._recycle_browser("active")

    @staticmethod
    def _build_page_options(
        render_options: dict[str, Any], *, pooled: bool
    ) -> dict[str, Any]:
        options = render_options.copy()
        options.pop("wait", None)
        options.pop("type", None)
        options.pop("quality", None)
        options.pop("scale", None)
        options.pop("screenshot_scale", None)
        options.pop("screenshot_timeout", None)
        options.pop("full_page", None)
        options.pop("clip_selector", None)
        options.pop("clip_padding", None)
        options.pop("disable_animations", None)
        if pooled:
            options.pop("base_url", None)
        return options

    @staticmethod
    def _build_screenshot_options(render_options: dict[str, Any]) -> dict[str, Any]:
        scale = render_options.get("screenshot_scale", render_options.get("scale"))
        if scale not in ("css", "device"):
            scale = None
        return {
            "full_page": bool(render_options.get("full_page", True)),
            "type": render_options.get("type", "png"),
            "quality": render_options.get("quality"),
            "scale": scale,
            "timeout": render_options.get("screenshot_timeout", 30_000),
        }

    @staticmethod
    def _get_wait_timeout(render_options: dict[str, Any]) -> int:
        wait = render_options.get("wait", 0)
        if isinstance(wait, int):
            return max(wait, 0)
        return 0

    @staticmethod
    def _coerce_non_negative_int(value: Any, default: int = 0) -> int:
        try:
            value_int = int(value)
        except (TypeError, ValueError):
            return default
        return value_int if value_int >= 0 else default

    @classmethod
    def _should_use_context_pool(cls, render_options: dict[str, Any]) -> bool:
        for key in cls._POOL_UNSAFE_OPTION_KEYS:
            if key in render_options:
                return False
        return True

    async def _render_with_page(
        self,
        page: Any,
        html: str,
        template_path: str,
        render_options: dict[str, Any],
    ) -> bytes:
        if self._debug_console_log:
            page.on("console", lambda msg: logger.debug(f"浏览器控制台: {msg.text}"))
        await page.goto(template_path, wait_until="domcontentloaded")
        await page.set_content(html, wait_until=self._SET_CONTENT_WAIT_UNTIL)
        if bool(render_options.get("disable_animations", False)):
            await self._disable_page_animations(page)
        await self._wait_for_visual_stability(page)
        if wait_ms := self._get_wait_timeout(render_options):
            await page.wait_for_timeout(wait_ms)
        screenshot_options = self._build_screenshot_options(render_options)
        clip_selector = render_options.get("clip_selector")
        if isinstance(clip_selector, str) and clip_selector.strip():
            if image_bytes := await self._capture_by_selector(
                page,
                selector=clip_selector.strip(),
                screenshot_options=screenshot_options,
                clip_padding=self._coerce_non_negative_int(
                    render_options.get("clip_padding"),
                    self._CLIP_PADDING_DEFAULT,
                ),
            ):
                return image_bytes
        await self._optimize_full_page_capture(page, screenshot_options)
        return await page.screenshot(**screenshot_options)

    async def _disable_page_animations(self, page: Any) -> None:
        with contextlib.suppress(Exception):
            await page.add_style_tag(content=self._DISABLE_ANIMATIONS_STYLE)

    async def _capture_by_selector(
        self,
        page: Any,
        selector: str,
        screenshot_options: dict[str, Any],
        clip_padding: int,
    ) -> bytes | None:
        element = await page.query_selector(selector)
        if element is None:
            return None

        element_screenshot_options = {
            "type": screenshot_options.get("type", "png"),
            "quality": screenshot_options.get("quality"),
            "timeout": screenshot_options.get("timeout", 30_000),
        }
        with contextlib.suppress(Exception):
            box = await element.bounding_box()
            if box and clip_padding > 0:
                viewport = page.viewport_size or {}
                width = int(viewport.get("width") or 0)
                if width > 0:
                    target_height = int(box["y"] + box["height"] + clip_padding)
                    current_height = int(viewport.get("height") or 0)
                    if target_height > current_height:
                        await page.set_viewport_size(
                            {
                                "width": width,
                                "height": min(
                                    target_height,
                                    self._FULL_PAGE_VIEWPORT_MAX_HEIGHT,
                                ),
                            }
                        )

        if clip_padding <= 0:
            return await element.screenshot(**element_screenshot_options)

        with contextlib.suppress(Exception):
            clip_box = await element.bounding_box()
            if clip_box is None:
                return await element.screenshot(**element_screenshot_options)
            clip = {
                "x": max(clip_box["x"] - clip_padding, 0),
                "y": max(clip_box["y"] - clip_padding, 0),
                "width": clip_box["width"] + clip_padding * 2,
                "height": clip_box["height"] + clip_padding * 2,
            }
            page_options = {
                "type": screenshot_options.get("type", "png"),
                "quality": screenshot_options.get("quality"),
                "timeout": screenshot_options.get("timeout", 30_000),
                "clip": clip,
            }
            return await page.screenshot(**page_options)

        return await element.screenshot(**element_screenshot_options)

    async def _wait_for_visual_stability(self, page: Any) -> None:
        with contextlib.suppress(Exception):
            await page.wait_for_function(
                "() => document.readyState === 'complete'",
                timeout=self._READY_STATE_TIMEOUT_MS,
            )

        with contextlib.suppress(Exception):
            await page.wait_for_function(
                "() => Array.from(document.images || []).every(img => img.complete)",
                timeout=self._IMAGE_READY_TIMEOUT_MS,
            )

        with contextlib.suppress(Exception):
            await page.evaluate(
                """
                async (timeoutMs) => {
                    if (!document.fonts || !document.fonts.ready) return;
                    await Promise.race([
                        document.fonts.ready,
                        new Promise(resolve => setTimeout(resolve, timeoutMs)),
                    ]);
                }
                """,
                self._FONT_READY_TIMEOUT_MS,
            )

    async def _optimize_full_page_capture(
        self, page: Any, screenshot_options: dict[str, Any]
    ) -> None:
        if not bool(screenshot_options.get("full_page")):
            return

        viewport = page.viewport_size or {}
        width = viewport.get("width")
        height = viewport.get("height")
        if (
            not isinstance(width, int)
            or width <= 0
            or not isinstance(height, int)
            or height <= 0
        ):
            return

        with contextlib.suppress(Exception):
            content_size = await page.evaluate(
                """
                () => {
                    const body = document.body;
                    const doc = document.documentElement;
                    const bodyWidth = body ? Math.max(
                        body.scrollWidth,
                        body.offsetWidth,
                        body.clientWidth
                    ) : 0;
                    const bodyHeight = body ? Math.max(
                        body.scrollHeight,
                        body.offsetHeight,
                        body.clientHeight
                    ) : 0;
                    const docWidth = doc ? Math.max(
                        doc.scrollWidth,
                        doc.offsetWidth,
                        doc.clientWidth
                    ) : 0;
                    const docHeight = doc ? Math.max(
                        doc.scrollHeight,
                        doc.offsetHeight,
                        doc.clientHeight
                    ) : 0;
                    return {
                        width: Math.ceil(Math.max(bodyWidth, docWidth, 10)),
                        height: Math.ceil(Math.max(bodyHeight, docHeight, 10)),
                    };
                }
                """
            )
            if not isinstance(content_size, dict):
                return

            content_width = content_size.get("width")
            content_height = content_size.get("height")
            if not isinstance(content_width, int) or not isinstance(
                content_height, int
            ):
                return
            if (
                content_width < 10
                or content_height < 10
                or content_width > self._FULL_PAGE_VIEWPORT_MAX_WIDTH
                or content_height > self._FULL_PAGE_VIEWPORT_MAX_HEIGHT
            ):
                return

            target_width = max(width, content_width)
            target_height = max(height, content_height)
            await page.set_viewport_size(
                {"width": target_width, "height": target_height}
            )
            screenshot_options["full_page"] = False

    async def _render_with_oneoff_page(
        self,
        html: str,
        template_path: str,
        render_options: dict[str, Any],
    ) -> bytes:
        browser = await _get_browser_instance()
        page_options = self._build_page_options(render_options, pooled=False)
        page = await browser.new_page(**page_options)
        try:
            return await self._render_with_page(
                page, html, template_path, render_options
            )
        finally:
            with contextlib.suppress(Exception):
                await page.close()

    async def _acquire_context(self) -> Any:
        try:
            return self._context_pool.get_nowait()
        except asyncio.QueueEmpty:
            pass

        async with self._state_lock:
            if len(self._all_contexts) < self._CONTEXT_POOL_SIZE:
                create_new = True
            else:
                create_new = False

        if create_new:
            browser = await _get_browser_instance()
            context = await browser.new_context(
                viewport={"width": 800, "height": 10},
                device_scale_factor=2,
            )
            async with self._state_lock:
                self._all_contexts.add(context)
            return context
        return await self._context_pool.get()

    async def _release_context(self, context: Any, broken: bool = False) -> None:
        if broken:
            await self._discard_context(context)
            return

        async with self._state_lock:
            if self._closing:
                broken = True
            elif context not in self._all_contexts:
                broken = True
            else:
                self._context_pool.put_nowait(context)
                return

        if broken:
            await self._discard_context(context)

    async def _discard_context(self, context: Any) -> None:
        async with self._state_lock:
            existed = context in self._all_contexts
            if existed:
                self._all_contexts.remove(context)
        if existed:
            with contextlib.suppress(Exception):
                await context.close()

    async def _dispose_context_pool(self) -> None:
        async with self._state_lock:
            contexts = list(self._all_contexts)
            self._all_contexts.clear()
            while True:
                try:
                    self._context_pool.get_nowait()
                except asyncio.QueueEmpty:
                    break

        for context in contexts:
            with contextlib.suppress(Exception):
                await context.close()

    async def _render_with_context_pool(
        self,
        html: str,
        template_path: str,
        render_options: dict[str, Any],
    ) -> bytes:
        context = await self._acquire_context()
        page = None
        broken = False
        try:
            page = await context.new_page()
            page_options = self._build_page_options(render_options, pooled=True)
            viewport = page_options.get("viewport")
            if isinstance(viewport, dict):
                width = viewport.get("width")
                height = viewport.get("height")
                if isinstance(width, int) and isinstance(height, int):
                    await page.set_viewport_size({"width": width, "height": height})
            return await self._render_with_page(
                page, html, template_path, render_options
            )
        except Exception:
            broken = True
            raise
        finally:
            if page is not None:
                with contextlib.suppress(Exception):
                    await page.close()
            await self._release_context(context, broken=broken)

    async def _render_html(
        self,
        html: str,
        template_path: str,
        render_options: dict[str, Any],
    ) -> bytes:
        if self._should_use_context_pool(render_options):
            return await self._render_with_context_pool(
                html, template_path, render_options
            )
        return await self._render_with_oneoff_page(html, template_path, render_options)

    async def _recycle_browser(self, reason: str) -> None:
        async with self._recycle_lock:
            try:
                await self._dispose_context_pool()
                await _shutdown_browser_instance()
                current_rss = self._get_total_rss()
                if current_rss is not None:
                    self._update_rss_baseline_nolock(current_rss)
                await self._prewarm_browser_and_pool()
                logger.debug(
                    f"截图引擎触发回收({reason})，已重建浏览器实例。",
                    "PlaywrightEngine",
                )
            except Exception as e:
                logger.warning("浏览器实例重建失败。", "PlaywrightEngine", e=e)

    async def _prewarm_browser_and_pool(self) -> None:
        if self._closing:
            return
        try:
            browser = await _get_browser_instance()
        except Exception as e:
            logger.warning("截图引擎浏览器预热失败。", "PlaywrightEngine", e=e)
            return

        for _ in range(self._PREWARM_CONTEXT_COUNT):
            async with self._state_lock:
                if self._closing:
                    return
                if len(self._all_contexts) >= self._CONTEXT_POOL_SIZE:
                    return
                if self._context_pool.qsize() >= self._PREWARM_CONTEXT_COUNT:
                    return

            context = None
            try:
                context = await browser.new_context(
                    viewport={"width": 800, "height": 10},
                    device_scale_factor=2,
                )
                page = await context.new_page()
                await page.goto("about:blank", wait_until="domcontentloaded")
                await page.set_content(
                    "<html><body></body></html>",
                    wait_until="domcontentloaded",
                )
                await page.close()
            except Exception as e:
                logger.warning("截图引擎上下文预热失败。", "PlaywrightEngine", e=e)
                if context is not None:
                    with contextlib.suppress(Exception):
                        await context.close()
                return

            async with self._state_lock:
                if self._closing:
                    with contextlib.suppress(Exception):
                        await context.close()
                    return
                if context in self._all_contexts:
                    continue
                self._all_contexts.add(context)
                self._context_pool.put_nowait(context)

    async def _idle_recycle_loop(self) -> None:
        while True:
            await asyncio.sleep(self._IDLE_CHECK_INTERVAL_SECONDS)
            should_recycle = False
            async with self._state_lock:
                if self._closing:
                    return
                now = time.monotonic()
                if self._active_renders > 0:
                    continue
                if now - self._last_recycle_at < self._RECYCLE_COOLDOWN_SECONDS:
                    continue
                idle_for = now - self._last_render_finished_at
                if idle_for < self._IDLE_RECYCLE_SECONDS:
                    continue
                current_rss = self._get_total_rss()
                if current_rss is None:
                    continue
                threshold = self._get_dynamic_threshold_nolock(current_rss)
                if current_rss >= threshold:
                    self._last_recycle_at = now
                    should_recycle = True
            if should_recycle:
                await self._recycle_browser("idle")

    async def _render_and_store_result(
        self,
        key: str,
        html: str,
        base_url_for_browser: str,
        render_options: dict[str, Any],
    ) -> bytes:
        async with self._render_semaphore:
            await self._on_render_begin()
            try:
                result = await self._render_html(
                    html,
                    base_url_for_browser,
                    render_options,
                )
            finally:
                await self._on_render_end()

        async with self._state_lock:
            now = time.monotonic()
            self._recent_results[key] = (
                now + self._RECENT_RESULT_TTL_SECONDS,
                result,
            )
            self._recent_results.move_to_end(key)
            self._cleanup_recent_results_nolock(now)
        return result

    async def render(self, html: str, base_url_path: Path, **render_options) -> bytes:
        base_url_for_browser = self._normalize_base_url(base_url_path)

        final_render_options = {
            "viewport": {"width": 800, "height": 10},
            **render_options,
            "base_url": base_url_for_browser,
        }

        dedupe_key = self._build_render_key(
            html,
            base_url_for_browser,
            final_render_options,
        )

        owner = False
        async with self._state_lock:
            now = time.monotonic()
            self._cleanup_recent_results_nolock(now)
            if cached_result := self._get_recent_result_nolock(dedupe_key, now):
                return cached_result

            task = self._inflight_tasks.get(dedupe_key)
            if task is None:
                task = asyncio.create_task(
                    self._render_and_store_result(
                        dedupe_key,
                        html,
                        base_url_for_browser,
                        final_render_options,
                    )
                )
                self._inflight_tasks[dedupe_key] = task
                owner = True

        try:
            return await task
        finally:
            if owner:
                async with self._state_lock:
                    if self._inflight_tasks.get(dedupe_key) is task:
                        self._inflight_tasks.pop(dedupe_key, None)


class EngineManager:
    """
    引擎管理器，负责加载和提供具体的截图引擎实例。
    未来可在此处根据 Config 读取不同的驱动配置。
    """

    def __init__(self):
        self._engine_class: type[BaseScreenshotEngine] = PlaywrightEngine
        self._instance: BaseScreenshotEngine | None = None

    async def get_engine(self) -> BaseScreenshotEngine:
        if not self._instance:
            self._instance = self._engine_class()
            await self._instance.initialize()
        return self._instance

    async def close(self):
        if self._instance:
            await self._instance.close()
            self._instance = None


engine_manager = EngineManager()