#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
图片转PDF脚本
将一张或多张图片转换为A4尺寸的PDF文件
使用Pillow处理图片，避免FPDF处理大PNG时的内存溢出问题
"""

import sys
import os
import io

if sys.platform == 'win32':
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')

from PIL import Image
Image.MAX_IMAGE_PIXELS = 200_000_000  # 允许大图，避免DecompressionBomb警告

# A4 at 300dpi
A4_WIDTH_PX = 2480
A4_HEIGHT_PX = 3508
A4_WIDTH_MM = 210
A4_HEIGHT_MM = 297
MARGIN_MM = 10

MAX_PIXELS = A4_WIDTH_PX * A4_HEIGHT_PX  # cap image to A4@300dpi


def image_to_pdf(image_paths: list, output_pdf: str, dpi: int = 300, quality: int = 90) -> bool:
    # When exactly 2 images (e.g. ID card front+back), place them on the same A4 page
    if len(image_paths) == 2:
        return two_images_one_page(image_paths, output_pdf, dpi)

    pages = []

    for img_path in image_paths:
        try:
            img = Image.open(img_path)
        except Exception as e:
            print(f"ERROR: Cannot open {img_path}: {e}", file=sys.stderr)
            continue

        img = _to_rgb(img)

        # Scale down if larger than A4 at target dpi
        target_w = int(A4_WIDTH_MM * dpi / 25.4)
        target_h = int(A4_HEIGHT_MM * dpi / 25.4)
        avail_w = int((A4_WIDTH_MM - 2 * MARGIN_MM) * dpi / 25.4)
        avail_h = int((A4_HEIGHT_MM - 2 * MARGIN_MM) * dpi / 25.4)

        img = _fit_image(img, avail_w, avail_h)
        img = _cap_pixels(img)

        # Create A4 white canvas and center the image
        canvas = Image.new('RGB', (target_w, target_h), (255, 255, 255))
        x_offset = (target_w - img.size[0]) // 2
        y_offset = (target_h - img.size[1]) // 2
        canvas.paste(img, (x_offset, y_offset))
        pages.append(canvas)
        img.close()

    if not pages:
        print("ERROR: No valid images to convert", file=sys.stderr)
        return False

    # Save as PDF
    first = pages[0]
    rest = pages[1:] if len(pages) > 1 else []

    try:
        first.save(output_pdf, 'PDF', resolution=dpi, save_all=True, append_images=rest)
        return True
    except Exception as e:
        print(f"ERROR: Failed to save PDF: {e}", file=sys.stderr)
        return False


def two_images_one_page(image_paths: list, output_pdf: str, dpi: int = 300) -> bool:
    """Place 2 images (e.g. ID card front+back) on a single A4 page, stacked vertically.
    Scale each image to fill its available space (allow scale-up), matching FPDF behavior."""
    spacing_mm = 5
    target_w = int(A4_WIDTH_MM * dpi / 25.4)
    target_h = int(A4_HEIGHT_MM * dpi / 25.4)
    avail_w = int((A4_WIDTH_MM - 2 * MARGIN_MM) * dpi / 25.4)
    avail_h_per = int((A4_HEIGHT_MM - 2 * MARGIN_MM - spacing_mm) * dpi / 25.4 / 2)

    images = []
    for img_path in image_paths:
        try:
            img = Image.open(img_path)
        except Exception as e:
            print(f"ERROR: Cannot open {img_path}: {e}", file=sys.stderr)
            return False
        img = _to_rgb(img)

        # Always scale to fill available area (allow scale-up, same as FPDF)
        w, h = img.size
        ratio = min(avail_w / w, avail_h_per / h)
        new_w = int(w * ratio)
        new_h = int(h * ratio)
        img = img.resize((new_w, new_h), Image.LANCZOS)

        images.append(img)

    if not images:
        print("ERROR: No valid images to convert", file=sys.stderr)
        return False

    canvas = Image.new('RGB', (target_w, target_h), (255, 255, 255))
    margin_px = int(MARGIN_MM * dpi / 25.4)
    spacing_px = int(spacing_mm * dpi / 25.4)
    y = margin_px

    for img in images:
        x_offset = (target_w - img.size[0]) // 2
        canvas.paste(img, (x_offset, y))
        y += img.size[1] + spacing_px
        img.close()

    try:
        canvas.save(output_pdf, 'PDF', resolution=dpi)
        return True
    except Exception as e:
        print(f"ERROR: Failed to save PDF: {e}", file=sys.stderr)
        return False


def _to_rgb(img: Image.Image) -> Image.Image:
    if img.mode not in ('RGB', 'L'):
        bg = Image.new('RGB', img.size, (255, 255, 255))
        if img.mode in ('RGBA', 'LA', 'PA'):
            bg.paste(img, mask=img.split()[-1])
        else:
            bg.paste(img)
        return bg
    elif img.mode == 'L':
        return img.convert('RGB')
    return img


def _fit_image(img: Image.Image, avail_w: int, avail_h: int) -> Image.Image:
    """Scale image to fill available area proportionally (allow scale-up for small images)."""
    w, h = img.size
    ratio = min(avail_w / w, avail_h / h)
    new_w = int(w * ratio)
    new_h = int(h * ratio)
    return img.resize((new_w, new_h), Image.LANCZOS)


def _cap_pixels(img: Image.Image) -> Image.Image:
    w, h = img.size
    if w * h > MAX_PIXELS:
        ratio = (MAX_PIXELS / (w * h)) ** 0.5
        return img.resize((int(w * ratio), int(h * ratio)), Image.LANCZOS)
    return img


def main():
    if len(sys.argv) < 3:
        print("Usage: python image_to_pdf.py <output.pdf> <image1> [image2] ...")
        print("  dpi=300 and quality=90 are used by default")
        sys.exit(1)

    output_pdf = sys.argv[1]
    image_paths = sys.argv[2:]

    # Validate input files
    valid_paths = []
    for p in image_paths:
        if not os.path.exists(p):
            print(f"ERROR: File not found: {p}", file=sys.stderr)
            sys.exit(1)
        valid_paths.append(p)

    success = image_to_pdf(valid_paths, output_pdf)
    if success:
        print(f"OK: {output_pdf}")
        sys.exit(0)
    else:
        sys.exit(1)


if __name__ == '__main__':
    main()