gantt_export.py

"""
将甘特图完整 HTML 文档渲染为 PNG（headless Chrome），再转 JPG，并写入本地缓存目录供下载。
依赖：html2image（需本机已安装 Chrome/Chromium）、Pillow。
"""
import io
import os
from datetime import datetime
from typing import Optional, Tuple


def _build_export_html(html_document: str) -> str:
    """注入导出专用样式：关闭滚动裁切，避免截图只保留可视区域。"""
    export_css = """
<style>
html, body {
    margin: 0 !important;
    padding: 0 !important;
    background: #fff !important;
    width: max-content !important;
    display: inline-block !important;
}
.scroll-wrapper { overflow: visible !important; width: max-content !important; }
.schedule-container { width: max-content !important; }
</style>
"""
    if "</head>" in html_document:
        return html_document.replace("</head>", f"{export_css}\n</head>", 1)
    return f"{export_css}\n{html_document}"


def _trim_png_whitespace(png_bytes: bytes, threshold: int = 250, padding: int = 8) -> bytes:
    """裁切四周近白空白区域，让图像内容尽量铺满画面。"""
    from PIL import Image

    im = Image.open(io.BytesIO(png_bytes)).convert("RGB")
    gray = im.convert("L")
    mask = gray.point(lambda x: 255 if x < threshold else 0)
    bbox = mask.getbbox()
    if not bbox:
        return png_bytes

    left, top, right, bottom = bbox
    left = max(0, left - padding)
    top = max(0, top - padding)
    right = min(im.width, right + padding)
    bottom = min(im.height, bottom + padding)

    cropped = im.crop((left, top, right, bottom))
    out = io.BytesIO()
    cropped.save(out, format="PNG", optimize=True)
    return out.getvalue()


def html_document_to_png_bytes(html_document: str, width: int = 1680, height: int = 5000) -> bytes:
    """使用 html2image 将完整 HTML 页面渲染为高分辨率 PNG 字节。"""
    import tempfile

    from html2image import Html2Image

    prepared_html = _build_export_html(html_document)

    with tempfile.TemporaryDirectory() as td:
        hti = Html2Image(
            output_path=td,
            size=(width, height),
            custom_flags=[
                "--hide-scrollbars",
                "--force-device-scale-factor=3",
                "--default-background-color=ffffff",
            ],
        )
        hti.screenshot(html_str=prepared_html, save_as="gantt.png")
        path = os.path.join(td, "gantt.png")
        with open(path, "rb") as f:
            raw = f.read()
    return _trim_png_whitespace(raw)


def png_bytes_to_jpeg_bytes(png_bytes: bytes, quality: int = 95) -> bytes:
    from PIL import Image

    im = Image.open(io.BytesIO(png_bytes)).convert("RGB")
    buf = io.BytesIO()
    im.save(buf, format="JPEG", quality=quality, optimize=True, subsampling=0)
    return buf.getvalue()


def save_png_jpeg_to_cache(
    png_bytes: bytes,
    jpg_bytes: bytes,
    base_name: str,
    cache_dir: Optional[str] = None,
) -> Tuple[str, str]:
    """后台保存 PNG/JPG 到 cinema_cache/gantt_exports/，返回文件路径。"""
    if cache_dir is None:
        cache_dir = os.path.join("cinema_cache", "gantt_exports")
    os.makedirs(cache_dir, exist_ok=True)
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    safe = "".join(c if c.isalnum() or c in "-_" else "_" for c in base_name)[:80]
    png_path = os.path.join(cache_dir, f"{safe}_{ts}.png")
    jpg_path = os.path.join(cache_dir, f"{safe}_{ts}.jpg")
    with open(png_path, "wb") as f:
        f.write(png_bytes)
    with open(jpg_path, "wb") as f:
        f.write(jpg_bytes)
    return png_path, jpg_path


def build_png_and_jpeg_for_gantt(html_document: str, num_hall_rows: int) -> Tuple[bytes, bytes]:
    """
    根据影厅行数估算截图尺寸，生成更清晰且内容占满画面的 PNG/JPG 字节。
    """
    nrows = max(1, num_hall_rows)
    width = 1680
    height = min(22000, max(2200, 360 + nrows * 120))
    png = html_document_to_png_bytes(html_document, width=width, height=height)
    jpg = png_bytes_to_jpeg_bytes(png)
    return png, jpg