#!/usr/bin/env python3
"""Render selected document pages, slides, sheets, ranges, or images for multimodal review."""

from __future__ import annotations

import json
import math
import subprocess
import sys
import tempfile
import time
from pathlib import Path

from PIL import Image, ImageDraw, ImageFont, ImageOps

WORKSPACE = Path("/workspace").resolve()
ALLOWED_ROOTS = tuple((WORKSPACE / name).resolve() for name in ("inputs", "work", "references", "artifacts"))
MAX_PIXELS = 8_000_000
MAX_IMAGE_BYTES = 4 * 1024 * 1024
MAX_TOTAL_BYTES = 20 * 1024 * 1024
MAX_CELLS = 2_400
COMMAND_TIMEOUT_SECONDS = 45
TOTAL_TIMEOUT_SECONDS = 170


def safe_source(value: str) -> Path:
    """Resolve a workspace-relative source and reject traversal or unsupported roots."""
    path = (WORKSPACE / value).resolve()
    if not any(path == root or path.is_relative_to(root) for root in ALLOWED_ROOTS):
        raise ValueError("文档路径必须位于 inputs、work、references 或 artifacts")
    if not path.is_file():
        raise ValueError("文档不存在")
    return path


def run(command: list[str]) -> None:
    """Run a renderer command with bounded execution time and useful errors."""
    result = subprocess.run(command, capture_output=True, text=True, timeout=COMMAND_TIMEOUT_SECONDS, check=False)
    if result.returncode:
        detail = (result.stderr or result.stdout or "渲染命令失败").strip()
        raise RuntimeError(detail[-600:])


def office_pdf(source: Path, directory: Path) -> Path:
    """Convert an Office document to PDF with an isolated LibreOffice profile."""
    profile = directory / "lo-profile"
    run([
        "libreoffice", "--headless", f"-env:UserInstallation=file://{profile}",
        "--convert-to", "pdf", "--outdir", str(directory), str(source),
    ])
    target = directory / f"{source.stem}.pdf"
    if not target.is_file():
        raise RuntimeError("LibreOffice 未生成 PDF")
    return target


def render_pdf(source: Path, page: int, dpi: int, target: Path, fmt: str) -> None:
    """Render one one-based PDF page."""
    output_type = "jpeg" if fmt == "jpeg" else "png"
    stem = target.with_suffix("")
    run([
        "pdftoppm", "-f", str(page), "-l", str(page), "-singlefile", "-r", str(dpi),
        f"-{output_type}", str(source), str(stem),
    ])
    produced = stem.with_suffix(".jpg" if output_type == "jpeg" else ".png")
    if not produced.is_file():
        raise RuntimeError("未找到指定页，请检查页码")
    produced.replace(target.with_suffix(produced.suffix))


def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
    """Load a CJK font available in the runtime image."""
    names = [
        "/usr/share/fonts/opentype/noto/NotoSansCJK-Bold.ttc" if bold else
        "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc",
        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
    ]
    for name in names:
        if Path(name).is_file():
            return ImageFont.truetype(name, size)
    return ImageFont.load_default()


def cell_text(value: object) -> str:
    """Format a spreadsheet cell for visual review."""
    if value is None:
        return ""
    text = str(value).replace("\r", " ").replace("\n", " ")
    return text if len(text) <= 80 else f"{text[:79]}…"


def draw_grid(values: list[list[object]], target: Path, title: str) -> None:
    """Render a bounded spreadsheet range as a readable table image."""
    rows = len(values)
    columns = max((len(row) for row in values), default=0)
    if rows == 0 or columns == 0:
        raise ValueError("所选工作表范围为空")
    if rows * columns > MAX_CELLS:
        raise ValueError(f"所选范围包含 {rows * columns} 个单元格，请拆分为多个较小范围")

    body_font = font(18)
    title_font = font(20, True)
    widths = []
    for column in range(columns):
        longest = max((len(cell_text(row[column])) if column < len(row) else 0 for row in values), default=0)
        widths.append(min(240, max(72, 18 + longest * 11)))
    row_height = 38
    title_height = 48
    width = sum(widths) + 2
    height = title_height + rows * row_height + 2
    scale = min(1.0, math.sqrt(MAX_PIXELS / max(1, width * height)))
    if scale < 1:
        widths = [max(42, int(value * scale)) for value in widths]
        row_height = max(24, int(row_height * scale))
        title_height = max(34, int(title_height * scale))
        body_font = font(max(12, int(18 * scale)))
        title_font = font(max(14, int(20 * scale)), True)
        width = sum(widths) + 2
        height = title_height + rows * row_height + 2

    image = Image.new("RGB", (width, height), "white")
    painter = ImageDraw.Draw(image)
    painter.rectangle((0, 0, width - 1, title_height), fill="#f1f5f9")
    painter.text((12, 11), title, fill="#334155", font=title_font)
    y = title_height
    for row_index, row in enumerate(values):
        x = 1
        fill = "#f8fafc" if row_index % 2 else "#ffffff"
        for column, cell_width in enumerate(widths):
            painter.rectangle((x, y, x + cell_width, y + row_height), fill=fill, outline="#dbe3ed")
            painter.text((x + 7, y + 8), cell_text(row[column] if column < len(row) else None),
                         fill="#253247", font=body_font)
            x += cell_width
        y += row_height
    image.save(target, "PNG", optimize=True)


def render_sheet(source: Path, sheet: str | None, selected_range: str | None, target: Path) -> str:
    """Render a selected spreadsheet range and return the actual sheet/range label."""
    if source.suffix.lower() == ".xls":
        import pandas as pd
        book = pd.ExcelFile(source)
        sheet_name = sheet or book.sheet_names[0]
        frame = pd.read_excel(source, sheet_name=sheet_name, header=None)
        from openpyxl.utils.cell import range_boundaries
        if selected_range:
            min_col, min_row, max_col, max_row = range_boundaries(selected_range)
            frame = frame.iloc[min_row - 1:max_row, min_col - 1:max_col]
        else:
            frame = frame.iloc[:60, :20]
            selected_range = f"A1:{_column_name(max(1, frame.shape[1]))}{max(1, frame.shape[0])}"
        values = frame.where(frame.notna(), None).values.tolist()
    else:
        from openpyxl import load_workbook
        from openpyxl.utils.cell import range_boundaries
        book = load_workbook(source, read_only=True, data_only=False)
        sheet_name = sheet or book.sheetnames[0]
        if sheet_name not in book.sheetnames:
            raise ValueError(f"工作表不存在：{sheet_name}")
        worksheet = book[sheet_name]
        if selected_range:
            min_col, min_row, max_col, max_row = range_boundaries(selected_range)
        else:
            min_col, min_row = 1, 1
            max_col, max_row = min(20, worksheet.max_column), min(60, worksheet.max_row)
            selected_range = f"A1:{_column_name(max_col)}{max_row}"
        values = [
            list(row)
            for row in worksheet.iter_rows(
                min_row=min_row,
                max_row=max_row,
                min_col=min_col,
                max_col=max_col,
                values_only=True,
            )
        ]
        book.close()
    label = f"{sheet_name}!{selected_range}"
    draw_grid(values, target, label)
    return label


def _column_name(index: int) -> str:
    """Convert a one-based column index to an Excel column name."""
    value = ""
    while index:
        index, remainder = divmod(index - 1, 26)
        value = chr(65 + remainder) + value
    return value


def normalize_image(path: Path, target: Path) -> None:
    """Normalize an uploaded image and bound its dimensions."""
    with Image.open(path) as source:
        image = ImageOps.exif_transpose(source).convert("RGB")
        if image.width * image.height > MAX_PIXELS:
            scale = math.sqrt(MAX_PIXELS / (image.width * image.height))
            image = image.resize((max(1, int(image.width * scale)), max(1, int(image.height * scale))))
        image.save(target, "JPEG", quality=88, optimize=True)


def bound_image(path: Path) -> Path:
    """Bound image pixels and bytes, converting oversized PNG files to JPEG."""
    with Image.open(path) as source:
        width, height = source.size
        if width * height > MAX_PIXELS:
            scale = math.sqrt(MAX_PIXELS / (width * height))
            image = source.convert("RGB").resize((max(1, int(width * scale)), max(1, int(height * scale))))
            replacement = path.with_suffix(".jpg")
            image.save(replacement, "JPEG", quality=88, optimize=True)
            path.unlink(missing_ok=True)
            path = replacement
    if path.stat().st_size > MAX_IMAGE_BYTES and path.suffix.lower() != ".jpg":
        with Image.open(path) as source:
            replacement = path.with_suffix(".jpg")
            source.convert("RGB").save(replacement, "JPEG", quality=84, optimize=True)
        path.unlink(missing_ok=True)
        path = replacement
    if path.stat().st_size > MAX_IMAGE_BYTES:
        raise ValueError("渲染图片过大，请降低 DPI 或缩小工作表范围")
    return path


def render(view: dict, index: int, output: Path) -> dict:
    """Render one requested view and return persisted metadata."""
    source = safe_source(str(view.get("path", "")))
    page = max(1, int(view.get("page") or 1))
    dpi = min(220, max(72, int(view.get("dpi") or 144)))
    fmt = "jpeg" if str(view.get("format", "png")).lower() in {"jpg", "jpeg"} else "png"
    suffix = source.suffix.lower()
    target = output / f"render-{index}.{'jpg' if fmt == 'jpeg' else 'png'}"
    label = f"第 {page} 页"

    if suffix == ".pdf":
        render_pdf(source, page, dpi, target, fmt)
    elif suffix in {".pptx", ".docx"}:
        with tempfile.TemporaryDirectory(dir=output) as temporary:
            converted = office_pdf(source, Path(temporary))
            render_pdf(converted, page, dpi, target, fmt)
        label = f"第 {page} 张" if suffix == ".pptx" else f"第 {page} 页"
    elif suffix in {".xls", ".xlsx"}:
        target = target.with_suffix(".png")
        label = render_sheet(source, view.get("sheet"), view.get("range"), target)
    elif suffix in {".png", ".jpg", ".jpeg", ".webp"}:
        target = target.with_suffix(".jpg")
        normalize_image(source, target)
        label = "原图"
    else:
        raise ValueError(f"document_view 暂不支持 {suffix or '无扩展名'} 文件")

    target = bound_image(target)
    with Image.open(target) as image:
        width, height = image.size
    return {
        "index": index,
        "sourcePath": str(source.relative_to(WORKSPACE)),
        "path": str(target.relative_to(WORKSPACE)),
        "label": label,
        "mediaType": "image/jpeg" if target.suffix.lower() == ".jpg" else "image/png",
        "width": width,
        "height": height,
        "sizeBytes": target.stat().st_size,
    }


def main(request_path: Path, result_path: Path) -> int:
    """Process all requested views and persist partial successes plus errors."""
    request_path = request_path.resolve()
    result_path = result_path.resolve()
    request = json.loads(request_path.read_text(encoding="utf-8"))
    views = request.get("views")
    if not isinstance(views, list) or not views:
        raise ValueError("views 至少包含一个查看请求")
    output = result_path.parent
    output.mkdir(parents=True, exist_ok=True)
    started = time.monotonic()
    images: list[dict] = []
    errors: list[dict] = []
    total_bytes = 0
    for index, view in enumerate(views, 1):
        if time.monotonic() - started > TOTAL_TIMEOUT_SECONDS:
            errors.append({"index": index, "error": "本次渲染已达到时间上限，请减少页面后重试"})
            continue
        try:
            image = render(view, index, output)
            total_bytes += image["sizeBytes"]
            if total_bytes > MAX_TOTAL_BYTES:
                Path(WORKSPACE / image["path"]).unlink(missing_ok=True)
                total_bytes -= image["sizeBytes"]
                raise ValueError("本次渲染图片总量过大，请拆分调用")
            images.append(image)
        except Exception as exception:  # Each failed view is returned to the Agent as an observation.
            errors.append({"index": index, "path": str(view.get("path", "")), "error": str(exception)[:800]})
    result_path.write_text(json.dumps({"images": images, "errors": errors}, ensure_ascii=False), encoding="utf-8")
    return 0 if images else 2


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("usage: document_view.py request.json result.json", file=sys.stderr)
        raise SystemExit(2)
    try:
        raise SystemExit(main(Path(sys.argv[1]), Path(sys.argv[2])))
    except Exception as failure:
        Path(sys.argv[2]).write_text(
            json.dumps({"images": [], "errors": [{"index": 0, "error": str(failure)[:800]}]}, ensure_ascii=False),
            encoding="utf-8",
        )
        raise SystemExit(2)