#!/bin/python3
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from PyPDF2 import PdfReader, PdfWriter
from reportlab.lib.units import inch
import sys

def create_watermark(output_pdf, watermark_text, font_path, font_size=32, opacity=0.5, angle=45, spacing_scale=1.5, margin=0.5):
    """
    创建包含多行水印的 PDF 文件，水印平整且均匀分布。
    :param output_pdf: 输出 PDF 文件路径
    :param watermark_text: 水印文本
    :param font_path: 字体文件路径
    :param font_size: 字体大小
    :param opacity: 水印透明度
    :param angle: 水印旋转角度（默认 45 度）
    :param spacing_scale: 水印间隔的缩放比例
    :param margin: 页面边距（单位：英寸）
    """
    # 注册字体
    pdfmetrics.registerFont(TTFont('CustomFont', font_path))

    # 创建 Canvas 对象
    c = canvas.Canvas(output_pdf, pagesize=letter)
    width, height = letter

    # 设置字体和透明度
    c.setFont("CustomFont", font_size)
    c.setFillAlpha(opacity)

    # 计算水印文本的宽度和高度
    text_width = c.stringWidth(watermark_text, "CustomFont", font_size)
    text_height = font_size

    # 计算水印间距，确保不会与页面边缘重叠
    spacing = max(text_width, text_height) * spacing_scale  # 水印间距为文本宽度或高度的 1.5 倍

    # 计算水印数量，确保水印不会超出页面边缘
    num_copies_x = int((width - 2 * margin * inch) / spacing) + 2  # 水平方向水印数量
    num_copies_y = int((height - 2 * margin * inch) / spacing) + 2  # 垂直方向水印数量

    # 重新计算水印布局，避免顶部和底部被截断
    # 这里增加了对旋转角度的处理，考虑到旋转后的文本可能会超出页面范围
    for i in range(num_copies_x):
        for j in range(num_copies_y):
            x = margin * inch + i * spacing  # 水平起始位置
            y = margin * inch + j * spacing  # 垂直起始位置

            # 旋转坐标系并绘制文本
            c.saveState()
            c.translate(x, y)
            c.rotate(angle)  # 旋转指定角度
            c.drawString(0, 0, watermark_text)
            c.restoreState()

    # 保存 PDF
    c.save()
    print(f"水印 PDF 已生成：{output_pdf}")


def get_arg(index, default_value):
    """获取命令行参数，如果没有提供则返回默认值"""
    return sys.argv[index] if len(sys.argv) > index else default_value

if __name__ == "__main__":
    # 如果传递了命令行参数，则覆盖默认值
    # 获取命令行参数或使用默认值
    output_pdf = get_arg(1, "watermark-01.pdf")
    watermark_text = get_arg(2, "kylin--ceshi")

    create_watermark(
        output_pdf,
        watermark_text,
        font_path="/usr/share/fonts/kyfonts/STSONG.TTF",  # 替换为你的字体文件路径
        font_size=16,
        opacity=0.5,
        angle=45,  # 水印旋转角度
        spacing_scale=1.5,  # 水印间隔的缩放比例
        margin=1  # 页面边距，单位为英寸
    )