watermark_tool.py

#!/usr/bin/env python3
"""Add text watermarks to PDF files and images on Windows.

Examples:
  python watermark_tool.py input.pdf output.pdf --text "CONFIDENTIAL" --position tile
  python watermark_tool.py photo.jpg photo_marked.jpg --text "SAMPLE" --opacity 0.2
"""

from __future__ import annotations

import argparse
import io
from pathlib import Path
from typing import Iterable, List, Tuple

PDF_EXTENSIONS = {".pdf"}
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"}


def parse_color(value: str) -> Tuple[int, int, int]:
    text = value.strip().lstrip("#")
    if len(text) != 6:
        raise argparse.ArgumentTypeError("--color must be a hex value like #FF0000")
    try:
        red = int(text[0:2], 16)
        green = int(text[2:4], 16)
        blue = int(text[4:6], 16)
    except ValueError as exc:
        raise argparse.ArgumentTypeError("--color must be a valid hex value like #FF0000") from exc
    return red, green, blue


def parse_opacity(value: str) -> float:
    try:
        opacity = float(value)
    except ValueError as exc:
        raise argparse.ArgumentTypeError("--opacity must be a number between 0 and 1") from exc
    if not (0.0 <= opacity <= 1.0):
        raise argparse.ArgumentTypeError("--opacity must be between 0 and 1")
    return opacity


def detect_mode(input_path: Path, requested_mode: str) -> str:
    if requested_mode != "auto":
        return requested_mode

    suffix = input_path.suffix.lower()
    if suffix in PDF_EXTENSIONS:
        return "pdf"
    if suffix in IMAGE_EXTENSIONS:
        return "image"
    raise ValueError(
        "Could not detect file type from extension. Use --mode pdf or --mode image explicitly."
    )


def get_image_font(font_size: int, font_path: str | None):
    from PIL import ImageFont

    candidates: List[str] = []
    if font_path:
        candidates.append(font_path)

    candidates.extend(
        [
            r"C:\\Windows\\Fonts\\arial.ttf",
            r"C:\\Windows\\Fonts\\segoeui.ttf",
            r"C:\\Windows\\Fonts\\meiryo.ttc",
        ]
    )

    for candidate in candidates:
        path = Path(candidate)
        if path.exists():
            try:
                return ImageFont.truetype(str(path), font_size)
            except OSError:
                continue

    return ImageFont.load_default()


def image_text_size(text: str, font) -> Tuple[int, int]:
    from PIL import Image, ImageDraw

    dummy = Image.new("RGBA", (1, 1), (0, 0, 0, 0))
    drawer = ImageDraw.Draw(dummy)
    left, top, right, bottom = drawer.textbbox((0, 0), text, font=font)
    return max(1, right - left), max(1, bottom - top)


def draw_rotated_text(layer, text: str, center: Tuple[float, float], font, fill, angle: float) -> None:
    from PIL import Image, ImageDraw

    text_width, text_height = image_text_size(text, font)
    patch = Image.new("RGBA", (text_width + 20, text_height + 20), (0, 0, 0, 0))
    drawer = ImageDraw.Draw(patch)
    drawer.text((10, 10), text, font=font, fill=fill)

    rotated = patch.rotate(angle, expand=True)
    x = int(center[0] - rotated.width / 2)
    y = int(center[1] - rotated.height / 2)

    layer.paste(rotated, (x, y), rotated)


def single_position_centers(
    width: int,
    height: int,
    position: str,
    margin: int,
) -> List[Tuple[float, float]]:
    if position == "center":
        return [(width / 2, height / 2)]
    if position == "top-left":
        return [(margin, margin)]
    if position == "top-right":
        return [(width - margin, margin)]
    if position == "bottom-left":
        return [(margin, height - margin)]
    if position == "bottom-right":
        return [(width - margin, height - margin)]
    raise ValueError(f"Unsupported position: {position}")


def tile_centers(width: int, height: int, step_x: int, step_y: int) -> Iterable[Tuple[float, float]]:
    y = -step_y
    while y <= height + step_y:
        x = -step_x
        while x <= width + step_x:
            yield (x, y)
            x += step_x
        y += step_y


def add_watermark_to_image(
    input_path: Path,
    output_path: Path,
    text: str,
    font_size: int,
    opacity: float,
    color: Tuple[int, int, int],
    angle: float,
    position: str,
    margin: int,
    font_path: str | None,
) -> None:
    try:
        from PIL import Image
    except ModuleNotFoundError as exc:
        raise RuntimeError("Pillow is required for image watermarking. Install with pip install -r requirements.txt") from exc

    base = Image.open(input_path).convert("RGBA")
    overlay = Image.new("RGBA", base.size, (0, 0, 0, 0))

    font = get_image_font(font_size, font_path)
    alpha = int(opacity * 255)
    fill = (color[0], color[1], color[2], alpha)
    text_width, text_height = image_text_size(text, font)

    if position == "tile":
        step_x = max(int(text_width * 2.4), 180)
        step_y = max(int(text_height * 2.8), 140)
        for center in tile_centers(base.width, base.height, step_x, step_y):
            draw_rotated_text(overlay, text, center, font, fill, angle)
    else:
        centers = single_position_centers(base.width, base.height, position, margin)
        for center in centers:
            draw_rotated_text(overlay, text, center, font, fill, angle)

    merged = Image.alpha_composite(base, overlay)

    output_path.parent.mkdir(parents=True, exist_ok=True)

    if output_path.suffix.lower() in {".jpg", ".jpeg", ".bmp"}:
        merged.convert("RGB").save(output_path)
    else:
        merged.save(output_path)


def draw_pdf_single_watermark(
    canvas_obj,
    page_width: float,
    page_height: float,
    text: str,
    font_size: int,
    angle: float,
    position: str,
    margin: int,
) -> None:
    from reportlab.pdfbase import pdfmetrics

    text_width = pdfmetrics.stringWidth(text, "Helvetica", font_size)

    if position == "center":
        center_x, center_y = page_width / 2, page_height / 2
    elif position == "top-left":
        center_x, center_y = margin, page_height - margin
    elif position == "top-right":
        center_x, center_y = page_width - margin, page_height - margin
    elif position == "bottom-left":
        center_x, center_y = margin, margin
    elif position == "bottom-right":
        center_x, center_y = page_width - margin, margin
    else:
        raise ValueError(f"Unsupported position: {position}")

    canvas_obj.saveState()
    canvas_obj.translate(center_x, center_y)
    canvas_obj.rotate(angle)
    canvas_obj.drawString(-text_width / 2, -font_size / 2, text)
    canvas_obj.restoreState()


def draw_pdf_tiled_watermark(
    canvas_obj,
    page_width: float,
    page_height: float,
    text: str,
    font_size: int,
    angle: float,
    margin: int,
) -> None:
    from reportlab.pdfbase import pdfmetrics

    text_width = pdfmetrics.stringWidth(text, "Helvetica", font_size)
    step_x = max(text_width * 2.0, 180)
    step_y = max(font_size * 3.0, 120)

    canvas_obj.saveState()
    canvas_obj.translate(page_width / 2, page_height / 2)
    canvas_obj.rotate(angle)

    y = -page_height
    while y <= page_height:
        x = -page_width
        while x <= page_width:
            canvas_obj.drawString(x, y, text)
            x += step_x
        y += step_y

    canvas_obj.restoreState()


def add_watermark_to_pdf(
    input_path: Path,
    output_path: Path,
    text: str,
    font_size: int,
    opacity: float,
    color: Tuple[int, int, int],
    angle: float,
    position: str,
    margin: int,
) -> None:
    try:
        from pypdf import PdfReader, PdfWriter
    except ModuleNotFoundError as exc:
        raise RuntimeError("pypdf is required for PDF watermarking. Install with pip install -r requirements.txt") from exc

    try:
        from reportlab.lib.colors import Color
        from reportlab.pdfgen import canvas
    except ModuleNotFoundError as exc:
        raise RuntimeError("reportlab is required for PDF watermarking. Install with pip install -r requirements.txt") from exc

    reader = PdfReader(str(input_path))
    writer = PdfWriter()

    for page in reader.pages:
        width = float(page.mediabox.width)
        height = float(page.mediabox.height)

        packet = io.BytesIO()
        c = canvas.Canvas(packet, pagesize=(width, height))
        c.setFont("Helvetica", font_size)

        fill_color = Color(color[0] / 255.0, color[1] / 255.0, color[2] / 255.0, alpha=opacity)
        c.setFillColor(fill_color)

        if hasattr(c, "setFillAlpha"):
            c.setFillAlpha(opacity)

        if position == "tile":
            draw_pdf_tiled_watermark(c, width, height, text, font_size, angle, margin)
        else:
            draw_pdf_single_watermark(c, width, height, text, font_size, angle, position, margin)

        c.save()
        packet.seek(0)

        watermark_page = PdfReader(packet).pages[0]
        page.merge_page(watermark_page)
        writer.add_page(page)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    with output_path.open("wb") as output_file:
        writer.write(output_file)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Add text watermarks to PDF files or images.")
    parser.add_argument("input", type=Path, help="Input file path (.pdf, .jpg, .png, etc.)")
    parser.add_argument("output", type=Path, help="Output file path")
    parser.add_argument("--text", required=True, help="Watermark text")
    parser.add_argument("--mode", choices=["auto", "pdf", "image"], default="auto", help="How to process input")
    parser.add_argument("--font-size", type=int, default=48, help="Font size (default: 48)")
    parser.add_argument("--opacity", type=parse_opacity, default=0.2, help="Opacity from 0 to 1 (default: 0.2)")
    parser.add_argument("--color", type=parse_color, default=(220, 0, 0), help="Hex color (default: #DC0000)")
    parser.add_argument("--angle", type=float, default=35.0, help="Rotation angle in degrees (default: 35)")
    parser.add_argument(
        "--position",
        choices=["center", "tile", "top-left", "top-right", "bottom-left", "bottom-right"],
        default="tile",
        help="Watermark position (default: tile)",
    )
    parser.add_argument("--margin", type=int, default=36, help="Margin for corner positions (default: 36)")
    parser.add_argument("--font-path", default=None, help="Optional TTF/TTC font path for image watermarking")
    return parser


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()

    if not args.input.exists() or not args.input.is_file():
        parser.error(f"Input file not found: {args.input}")

    if args.font_size <= 0:
        parser.error("--font-size must be > 0")

    try:
        mode = detect_mode(args.input, args.mode)
    except ValueError as exc:
        parser.error(str(exc))

    try:
        if mode == "pdf":
            add_watermark_to_pdf(
                input_path=args.input,
                output_path=args.output,
                text=args.text,
                font_size=args.font_size,
                opacity=args.opacity,
                color=args.color,
                angle=args.angle,
                position=args.position,
                margin=args.margin,
            )
        else:
            add_watermark_to_image(
                input_path=args.input,
                output_path=args.output,
                text=args.text,
                font_size=args.font_size,
                opacity=args.opacity,
                color=args.color,
                angle=args.angle,
                position=args.position,
                margin=args.margin,
                font_path=args.font_path,
            )
    except Exception as exc:
        parser.error(str(exc))

    print(f"Saved: {args.output}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())