mytoolkit/bin/commands/pdf.py

"""PDF utilities."""

import shutil
from pathlib import Path

import click
from pypdf import PdfReader, PdfWriter
from PIL import Image

from bin.md_to_pdf import expand_bookmarks
from bin.utils import handle_errors, run_command


@click.group()
def pdf():
    """PDF manipulation commands."""
    pass


@pdf.command("compress")
@click.argument("files", nargs=-1, required=True)
@click.option(
    "--quality",
    "-q",
    type=click.Choice(["screen", "ebook", "printer", "prepress"]),
    default="screen",
    help="Compression quality level",
)
@click.option(
    "--bookmark-depth",
    type=int,
    default=1,
    help="Expand PDF bookmarks to this depth after compression (1=no expansion).",
)
@click.option("--dry-run", "-n", is_flag=True, help="Show what would be done")
@handle_errors
def compress(files, quality, bookmark_depth, dry_run):
    """Compress PDF files using ghostscript."""
    for pattern in files:
        for pdf_file in Path(".").glob(pattern):
            if pdf_file.suffix.lower() != ".pdf":
                continue

            output = pdf_file.with_suffix(".compressed.pdf")
            cmd = [
                "gs",
                "-sDEVICE=pdfwrite",
                "-dNOPAUSE",
                "-dQUIET",
                "-dBATCH",
                f"-dPDFSETTINGS=/{quality}",
                "-dCompatibilityLevel=1.4",
                f"-sOutputFile={output}",
                str(pdf_file),
            ]

            if dry_run:
                click.echo(f"Would compress: {pdf_file} -> {output}")
                continue

            click.echo(f"Compressing: {pdf_file}")
            run_command(cmd, check=True)
            output.replace(pdf_file)
            if bookmark_depth > 1:
                expand_bookmarks(pdf_file, bookmark_depth)
            click.echo(f"  Done: {pdf_file}")


@pdf.command("crop")
@click.argument("files", nargs=-1, required=True)
@click.option("--dry-run", "-n", is_flag=True, help="Show what would be done")
@handle_errors
def crop(files, dry_run):
    """Crop PDF files."""
    for pattern in files:
        for pdf_file in Path(".").glob(pattern):
            if pdf_file.suffix.lower() != ".pdf":
                continue

            if dry_run:
                click.echo(f"Would crop: {pdf_file}")
                continue

            click.echo(f"Cropping: {pdf_file}")
            run_command(["pdfcrop", str(pdf_file), str(pdf_file)], check=False)


@pdf.command("to-tiff")
@click.argument("files", nargs=-1, required=True)
@click.option("--density", "-d", default="300", help="DPI density")
@click.option("--dry-run", "-n", is_flag=True, help="Show what would be done")
@handle_errors
def to_tiff(files, density, dry_run):
    """Convert PDF to TIFF."""
    for pattern in files:
        for pdf_file in Path(".").glob(pattern):
            if pdf_file.suffix.lower() != ".pdf":
                continue
            output = pdf_file.with_suffix(".tiff")

            if dry_run:
                click.echo(f"Would convert: {pdf_file} -> {output}")
                continue

            click.echo(f"Converting: {pdf_file} -> {output}")
            run_command(
                ["convert", "-density", density, str(pdf_file), "-compress", "zip", str(output)],
                check=False,
            )


@pdf.command("merge")
@click.argument("files", nargs=-1, required=True)
@click.option("--output", "-o", required=True, help="Output PDF file")
@click.option(
    "--bookmark-depth",
    type=int,
    default=1,
    help="Expand PDF bookmarks to this depth after merge (1=no expansion).",
)
@click.option("--dry-run", "-n", is_flag=True, help="Show what would be done")
@handle_errors
def merge_pdfs(files, output, bookmark_depth, dry_run):
    """Merge PDFs and/or images into a single PDF."""
    writer = PdfWriter()
    processed = []

    for pattern in files:
        for path in sorted(Path(".").glob(pattern)):
            if path.suffix.lower() == ".pdf":
                if dry_run:
                    click.echo(f"Would add PDF: {path}")
                    continue
                reader = PdfReader(str(path))
                for page in reader.pages:
                    writer.add_page(page)
                processed.append(str(path))
            elif path.suffix.lower() in (".png", ".jpg", ".jpeg"):
                if dry_run:
                    click.echo(f"Would add image: {path}")
                    continue
                img = Image.open(path)
                if img.mode in ("RGBA", "P"):
                    img = img.convert("RGB")
                # Save to temp PDF
                temp_pdf = path.with_suffix(".temp.pdf")
                img.save(temp_pdf, "PDF", resolution=150.0)
                reader = PdfReader(str(temp_pdf))
                writer.add_page(reader.pages[0])
                temp_pdf.unlink()
                processed.append(str(path))

    if dry_run:
        click.echo(f"Would write merged PDF to: {output}")
        return

    if not processed:
        click.echo("No files found to merge")
        return

    with open(output, "wb") as f:
        writer.write(f)
    if bookmark_depth > 1:
        expand_bookmarks(Path(output), bookmark_depth)
    click.echo(f"Merged {len(processed)} files into {output}")


@pdf.command("bookmark")
@click.argument("input_file", type=click.Path(exists=True, dir_okay=False, path_type=Path))
@click.option(
    "--depth",
    type=int,
    required=True,
    help="Bookmark expansion depth. 1=top-level only, 2=expand to second level, etc.",
)
@click.option(
    "-o",
    "--output",
    type=click.Path(dir_okay=False, path_type=Path),
    help="Output PDF path. Defaults to overwriting the input file.",
)
@handle_errors
def bookmark(input_file, depth, output):
    """Set PDF bookmark outline expansion depth."""
    target = output if output else input_file
    if output and output.resolve() != input_file.resolve():
        shutil.copy2(str(input_file), str(target))
    expand_bookmarks(target, depth)
    click.echo(f"Bookmarks expanded to depth {depth}: {target}")