import fitz  # PyMuPDF
from PIL import Image
import cv2
import numpy as np
from pyzbar import pyzbar
import io


#pip install PyMuPDF Pillow opencv-python numpy pyzbar

def extract_qr_from_pdf(pdf_path, output_image_path='qrcode.png', page_num=0):
    """
    从PDF中识别并提取二维码图片
    
    参数:
        pdf_path: PDF文件路径
        output_image_path: 输出二维码图片的路径
        page_num: 要处理的页码(从0开始),默认为第一页
    
    返回:
        成功返回True和二维码数据,失败返回False和None
    """
    
    try:
        # 打开PDF文件
        pdf_document = fitz.open(pdf_path)
        
        # 获取指定页面
        if page_num >= len(pdf_document):
            print(f"页码超出范围,PDF只有{len(pdf_document)}页")
            return False, None
            
        page = pdf_document[page_num]
        
        # 将页面转换为图像(提高分辨率以便更好识别)
        mat = fitz.Matrix(3.0, 3.0)  # 放大3倍
        pix = page.get_pixmap(matrix=mat)
        
        # 转换为PIL Image
        img_data = pix.tobytes("png")
        img = Image.open(io.BytesIO(img_data))
        
        # 转换为OpenCV格式
        img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
        
        # 检测二维码
        qr_codes = pyzbar.decode(img_cv)
        
        if not qr_codes:
            print("未在PDF中检测到二维码")
            return False, None
        
        # 提取第一个检测到的二维码
        qr = qr_codes[0]
        
        # 获取二维码的位置
        x, y, w, h = qr.rect
        
        # 添加一些边距使截图更完整
        margin = 10
        x = max(0, x - margin)
        y = max(0, y - margin)
        w = w + 2 * margin
        h = h + 2 * margin
        
        # 截取二维码区域
        qr_image = img_cv[y:y+h, x:x+w]
        
        # 保存二维码图片
        cv2.imwrite(output_image_path, qr_image)
        
        # 获取二维码数据
        qr_data = qr.data.decode('utf-8')
        
        print(f"成功提取二维码!")
        print(f"二维码内容: {qr_data}")
        print(f"二维码图片已保存至: {output_image_path}")
        
        pdf_document.close()
        return True, qr_data
        
    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
        return False, None


def extract_all_qr_from_pdf(pdf_path, output_dir='qr_codes'):
    """
    从PDF的所有页面中提取所有二维码
    
    参数:
        pdf_path: PDF文件路径
        output_dir: 输出目录
    """
    import os
    
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    pdf_document = fitz.open(pdf_path)
    total_qr_count = 0
    
    for page_num in range(len(pdf_document)):
        print(f"\n处理第 {page_num + 1} 页...")
        
        page = pdf_document[page_num]
        mat = fitz.Matrix(3.0, 3.0)
        pix = page.get_pixmap(matrix=mat)
        
        img_data = pix.tobytes("png")
        img = Image.open(io.BytesIO(img_data))
        img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
        
        qr_codes = pyzbar.decode(img_cv)
        
        for idx, qr in enumerate(qr_codes):
            x, y, w, h = qr.rect
            margin = 10
            x = max(0, x - margin)
            y = max(0, y - margin)
            w = w + 2 * margin
            h = h + 2 * margin
            
            qr_image = img_cv[y:y+h, x:x+w]
            output_path = os.path.join(output_dir, f'page{page_num+1}_qr{idx+1}.png')
            cv2.imwrite(output_path, qr_image)
            
            qr_data = qr.data.decode('utf-8')
            print(f"  - 二维码 {idx+1}: {qr_data}")
            print(f"    保存至: {output_path}")
            total_qr_count += 1
    
    pdf_document.close()
    print(f"\n总共提取了 {total_qr_count} 个二维码")


# 使用示例
if __name__ == "__main__":
    import sys
    
    # 支持命令行调用
    if len(sys.argv) >= 3:
        pdf_path = sys.argv[1]
        output_path = sys.argv[2]
        page_num = int(sys.argv[3]) if len(sys.argv) > 3 else 0
        
        success, qr_data = extract_qr_from_pdf(pdf_path, output_path, page_num)
        
        # 返回状态码: 0=成功, 1=失败
        sys.exit(0 if success else 1)
    else:
        print("用法: python qrcode.py <pdf_path> <output_path> [page_num]")
        print("示例: python qrcode.py license.pdf qrcode.png 0")
        sys.exit(1)