from datetime import date
from reportlab.lib import units
from reportlab.pdfgen import canvas
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
pdfmetrics.registerFont(TTFont('song', 'simsun.ttc'))
import PyPDF2


class WaterMarkUtils:
    @staticmethod
    def create_wm_template(content,
                           width=210,
                           height=297,
                           angle=135,
                           text_stroke_color_rgb=(0, 0, 0),
                           text_fill_color_rgb=(255, 0, 0),
                           text_fill_alpha=0.1,
                           show_date=True,
                           filename='template'
                           ) -> str:
      # 生成水印模板
        file_path = f'{filename}.pdf'
        c = canvas.Canvas(f'{filename}.pdf', pagesize=(width * units.mm, height * units.mm))
        # 画布平移保证文字完整性
        print(units.mm)
        # c.translate(0.4 * width * units.mm, 0.4 * height * units.mm)
        c.rotate(angle)
        c.setStrokeColorRGB(*text_stroke_color_rgb)
        c.setFont('song', size=32)
        # 设置填充色
        c.setFillColorRGB(*text_fill_color_rgb)
        # 设置字体透明度
        c.setFillAlpha(text_fill_alpha)
        # 绘制字体内容
        if show_date:
            cur_date = date.today().__str__()
            content += f'({cur_date})'
        # c.drawString(0, 0, content)
        # c.drawString(50, 100, content)
        # c.drawString(100, 150, content)
        # c.drawString(150, 200, content)
        # c.drawString(200, 250, content)
        # for i in range(100,1500,100):
            # c.drawString(i*0.7, i, "{}".format(i))
        c.drawString(600, 1300, content)
        # c.drawText(content)
        c.save()
        return file_path

    @staticmethod
    def add_watermark(content, source_path, out_path=None,angle=45, show_date=True,):
        template_path = WaterMarkUtils.create_wm_template(content, angle=angle, show_date=show_date)

        def add_wm_for_page(wm_file, page_pdf):
            # 打开水印pdf文件
            pdf_reader = PyPDF2.PdfFileReader(wm_file)
            page_pdf.mergePage(pdf_reader.getPage(0))
            return page_pdf

        source_reader = PyPDF2.PdfFileReader(source_path)
        target_writer = PyPDF2.PdfFileWriter()
        for page in range(source_reader.numPages):
            vm_pdf = add_wm_for_page(template_path, source_reader.getPage(page))
            target_writer.addPage(vm_pdf)
        if not out_path:
            target_file_name = source_path.split('.')[0] + '_wm.pdf'
        else:
            target_file_name = out_path
        target_writer.write(open(target_file_name, 'wb'))

if __name__ == "__main__":
    # WaterMarkUtils.create_wm_template("水印测试")
    WaterMarkUtils.add_watermark("水印测试","test.pdf",out_path="output.pdf",angle=0)