#!/usr/bin/python3.6
# -*- coding: utf-8 -*-
# @Author  : lijunliang
# @Email   : lijunliang@topnet.net.cn
# @File    : watermark.py
# @Software: PyCharm

"""
pdf文件添加水印
"""

from PyPDF2 import PdfFileReader, PdfFileWriter


def add_watermark(pdf_file_in: str, pdf_file_out: str, watermark_file: str):
    """
    把水印添加到pdf中
    :param pdf_file_in: 输入文件
    :param pdf_file_out: 输出文件
    :param watermark_file: 水印文件
    :return:
    """
    try:
        pdf_output = PdfFileWriter()
        input_stream = open(pdf_file_in, 'rb')
        pdf_input = PdfFileReader(input_stream, strict=False)

        # 获取PDF文件的页数
        pageNum = pdf_input.getNumPages()

        # 读入水印pdf文件
        output_stream = open(watermark_file, 'rb')
        pdf_watermark = PdfFileReader(output_stream)

        # 给每一页打水印
        for i in range(pageNum):
            page = pdf_input.getPage(i)
            page.mergePage(pdf_watermark.getPage(0))
            page.compressContentStreams()  # 压缩内容
            pdf_output.addPage(page)
        save_stream = open(pdf_file_out, 'wb')
        pdf_output.write(save_stream)
        save_stream.close()
        input_stream.close()
        output_stream.close()
        return True
    except Exception as e:
        print(e)
        return False

# if __name__ == '__main__':
#     pdf_file_in = '../0.pdf'
#     pdf_file_water = '../data/mark.pdf'
#     pdf_file_out = '../data/out.pdf'
#     add_watermark(pdf_file_in, pdf_file_out, pdf_file_water)