import fitz # PyMuPDF
from PIL import Image
from pathlib import Path
from io import BytesIO
import qrcode


class PDFStampAdder:
    def __init__(self, pdf_path_or_bytes):
        """初始化PDF骑缝章添加器

        Args:
            pdf_path_or_bytes: PDF文件路径或PDF文件的bytes
        """
        if isinstance(pdf_path_or_bytes, (str, Path)):
            pdf_path = Path(pdf_path_or_bytes)
            self.doc = fitz.open(pdf_path)
        else:
            stream = BytesIO(pdf_path_or_bytes)
            self.doc = fitz.open(stream=stream, filetype="pdf")

    def _compress_image(self, img, quality=30):
        """压缩图片

        Args:
            img: PIL Image对象
            quality: 压缩质量,1-100,默认30

        Returns:
            BytesIO: 压缩后的图片数据
        """
        output = BytesIO()
        # 保持PNG的透明通道
        img.save(output, format='PNG', optimize=True, quality=quality)
        output.seek(0)
        return output

    def add_across_page_stamp_to_pdf(self, stamp_path, scale_height=100, page_gap=0):
        """添加骑缝章到PDF
        
        Args:
            stamp_path: 骑缝章图片路径
            scale_height: 骑缝章缩放后的高度
        """
        # 获取骑缝章尺寸和缩放比例
        stamp_img = Image.open(stamp_path)
        stamp_height = stamp_img.height
        scale_ratio = scale_height / stamp_height

        # 打开骑缝章并计算切片宽度
        page_count = self.doc.page_count - page_gap
        slice_width = stamp_img.width // page_count

        for page_num in range(page_count):
            # 裁剪骑缝章片段
            slice = stamp_img.crop((
                page_num * slice_width,  # 左边界
                0,  # 上边界
                (page_num + 1) * slice_width,  # 右边界
                stamp_height  # 下边界
            ))
            
            # 压缩图片片段
            compressed_slice = self._compress_image(slice)
            
            # 添加到PDF页面
            page = self.doc.load_page(page_num)
            scaled_width = slice_width * scale_ratio
            
            # 调整位置到页面右侧居中
            rect = fitz.Rect(
                page.rect.width - scaled_width,  # 左边界贴着页面右边
                page.rect.height / 2 - scale_height/2,  # 上边界居中
                page.rect.width,  # 右边界超出页面（骑缝效果）
                page.rect.height / 2 + scale_height/2  # 下边界
            )
            
            page.insert_image(rect, stream=compressed_slice)

    def add_stamp_to_pdf(self, stamp_path_or_bytes, scale_height, page_num, x, y):
        """添加普通印章到PDF指定页面的指定位置
        
        Args:
            stamp_path: 印章图片路径
            scale_height: 印章缩放后的高度
            page_num: 页码（从0开始）
            x: 印章左上角的x坐标
            y: 印章左上角的y坐标
        """
        # 获取印章尺寸和缩放比例
        if isinstance(stamp_path_or_bytes, (str, Path)):
            stamp_img = Image.open(stamp_path_or_bytes)
        else:
            stamp_img = Image.open(BytesIO(stamp_path_or_bytes))
        
        stamp_height = stamp_img.height
        stamp_width = stamp_img.width
        scale_ratio = scale_height / stamp_height
        scaled_width = stamp_width * scale_ratio
        
        # 压缩印章图片
        compressed_stamp = self._compress_image(stamp_img)
        
        # 添加到PDF页面
        page = self.doc.load_page(page_num)
        rect = fitz.Rect(
            x,  # 左边界
            y,  # 上边界
            x + scaled_width,  # 右边界
            y + scale_height  # 下边界
        )
        
        page.insert_image(rect, stream=compressed_stamp)

    def add_qrcode_to_pdf(self, text, size, page_num, x, y):
        """添加二维码到PDF指定页面的指定位置
        
        Args:
            text: 二维码内容
            size: 二维码大小(像素)
            page_num: 页码（从0开始）
            x: 二维码左上角的x坐标
            y: 二维码左上角的y坐标
        """
        # 生成二维码
        qr = qrcode.QRCode(version=1, box_size=3, border=0)
        qr.add_data(text)
        qr.make(fit=True)
        qr_img = qr.make_image(fill_color="black", back_color="white")
        
        # 调整二维码大小
        qr_img = qr_img.resize((size, size), Image.LANCZOS)
        
        # 压缩二维码图片
        compressed_qr = self._compress_image(qr_img)
        
        # 添加到PDF页面
        page = self.doc.load_page(page_num)
        rect = fitz.Rect(
            x,  # 左边界
            y,  # 上边界
            x + size,  # 右边界
            y + size  # 下边界
        )
        
        page.insert_image(rect, stream=compressed_qr)

    def add_text_to_pdf(self, text, page_num, x, y, font_size=20):
        """添加文字到PDF指定页面的指定位置

        Args:
            text: 要添加的文字内容
            font_size: 字体大小，默认20
            page_num: 页码（从0开始）
            x: 文字左上角的x坐标
            y: 文字左上角的y坐标
            color: 文字颜色，RGB元组，默认黑色(0,0,0)
        """
        page = self.doc.load_page(page_num)
        
        # 创建文本写入器
        tw = fitz.TextWriter(page.rect)
        tw.append((x, y), text, fontsize=font_size)
        
        # 将文本写入页面
        tw.write_text(page)

    def add_text_around_text(self, text, add_text, page_num, offset_x=0, offset_y=0, font_size=20):
        """在PDF中搜索指定文本并在其周围添加文字
        
        Args:
            text: 要搜索的文本
            add_text: 要添加的文字内容
            page_num: 页码（从0开始）
            offset_x: 相对文字的x坐标偏移量
            offset_y: 相对文字的y坐标偏移量
            font_size: 字体大小，默认20
        """
        page = self.doc.load_page(page_num)
        text_instances = page.search_for(text)
        
        for inst in text_instances:
            # 计算添加文字的位置
            x = inst.x0 + offset_x
            y = inst.y0 + offset_y
            
            # 添加文字
            self.add_text_to_pdf(add_text, page_num, x, y, font_size)
            break

        print(f"Text '{add_text}' added around the text '{text}'.")

    def add_stamp_around_text(self, text, stamp_path_or_bytes, scale_height=30, offset_x=0, offset_y=0):
        """在PDF中搜索指定文本并在其周围添加印章
        
        Args:
            text: 要搜索的文本
            stamp_path_or_bytes: 印章图片路径或bytes
            scale_height: 印章缩放后的高度
            offset_x: 相对文字的x坐标偏移量
            offset_y: 相对文字的y坐标偏移量
        """
        # 获取印章尺寸和缩放比例
        if isinstance(stamp_path_or_bytes, (str, Path)):
            stamp_img = Image.open(stamp_path_or_bytes)
        else:
            stamp_img = Image.open(BytesIO(stamp_path_or_bytes))
        
        stamp_height = stamp_img.height
        stamp_width = stamp_img.width
        scale_ratio = scale_height / stamp_height
        scaled_width = stamp_width * scale_ratio
        
        # 搜索文本并添加印章
        for page_num in range(self.doc.page_count):
            page = self.doc.load_page(page_num)
            text_instances = page.search_for(text)
            for inst in text_instances:
                # 计算印章位置，考虑相对坐标偏移
                rect = fitz.Rect(
                    inst.x0 - scaled_width / 2 + offset_x,  # 左边界
                    inst.y0 - scale_height / 2 + offset_y,  # 上边界
                    inst.x0 + scaled_width / 2 + offset_x,  # 右边界
                    inst.y0 + scale_height / 2 + offset_y   # 下边界
                )
                
                # 压缩印章图片
                compressed_stamp = self._compress_image(stamp_img)
                
                # 添加印章到PDF页面
                page.insert_image(rect, stream=compressed_stamp)
                break

        print(f"Stamps added around the text '{text}'.")

    def save(self, output_path):
        """保存添加了骑缝章的PDF到指定路径
        
        Args:
            output_path: 输出PDF路径
        """
        # 使用压缩参数保存PDF
        self.doc.save(output_path, 
                     garbage=4,  # 最大垃圾回收
                     deflate=True,  # 使用deflate压缩
                     clean=True)  # 清理和优化PDF
        self.doc.close()
        print(f"Riding stamp has been added to {output_path}")

    @property
    def bytes(self):
        """将PDF文档转换为bytes

        Returns:
            bytes: PDF文档的二进制内容
        """
        buffer = BytesIO()
        self.doc.save(buffer,
                     garbage=4,
                     deflate=True,
                     clean=True)
        return buffer.getvalue()

