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


def extract_qr_from_pdf_opencv(pdf_path, output_image_path='qrcode.png', page_num=0):
    """
    使用OpenCV检测二维码（不依赖pyzbar）
    
    参数:
        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)
        
        # 使用OpenCV的QRCodeDetector
        qr_detector = cv2.QRCodeDetector()
        
        # 检测和解码二维码
        data, points, _ = qr_detector.detectAndDecode(img_cv)
        
        if data and points is not None:
            # 获取二维码的边界框
            points = points[0].astype(int)
            x_min = int(np.min(points[:, 0]))
            y_min = int(np.min(points[:, 1]))
            x_max = int(np.max(points[:, 0]))
            y_max = int(np.max(points[:, 1]))
            
            # 添加一些边距
            margin = 20
            x_min = max(0, x_min - margin)
            y_min = max(0, y_min - margin)
            x_max = min(img_cv.shape[1], x_max + margin)
            y_max = min(img_cv.shape[0], y_max + margin)
            
            # 截取二维码区域
            qr_image = img_cv[y_min:y_max, x_min:x_max]
            
            # 保存二维码图片
            cv2.imwrite(output_image_path, qr_image)
            
            print(f"成功提取二维码!")
            print(f"二维码内容: {data}")
            print(f"二维码图片已保存至: {output_image_path}")
            
            pdf_document.close()
            return True, data
        else:
            print("未在PDF中检测到二维码")
            pdf_document.close()
            return False, None
        
    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
        return False, None


def extract_all_qr_from_pdf_opencv(pdf_path, output_dir='qr_codes'):
    """
    从PDF的所有页面中提取所有二维码（使用OpenCV）
    """
    import os
    
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    pdf_document = fitz.open(pdf_path)
    total_qr_count = 0
    qr_detector = cv2.QRCodeDetector()
    
    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)
        
        # 检测多个二维码
        success, decoded_info, points, _ = qr_detector.detectAndDecodeMulti(img_cv)
        
        if success:
            for idx, (data, point_set) in enumerate(zip(decoded_info, points)):
                if data:  # 确保解码成功
                    points_int = point_set.astype(int)
                    x_min = int(np.min(points_int[:, 0]))
                    y_min = int(np.min(points_int[:, 1]))
                    x_max = int(np.max(points_int[:, 0]))
                    y_max = int(np.max(points_int[:, 1]))
                    
                    margin = 20
                    x_min = max(0, x_min - margin)
                    y_min = max(0, y_min - margin)
                    x_max = min(img_cv.shape[1], x_max + margin)
                    y_max = min(img_cv.shape[0], y_max + margin)
                    
                    qr_image = img_cv[y_min:y_max, x_min:x_max]
                    output_path = os.path.join(output_dir, f'page{page_num+1}_qr{idx+1}.png')
                    cv2.imwrite(output_path, qr_image)
                    
                    print(f"  - 二维码 {idx+1}: {data}")
                    print(f"    保存至: {output_path}")
                    total_qr_count += 1
    
    pdf_document.close()
    print(f"\n总共提取了 {total_qr_count} 个二维码")


# 使用示例
if __name__ == "__main__":
    # 支持命令行调用
    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_opencv(pdf_path, output_path, page_num)
        
        # 返回状态码: 0=成功, 1=失败
        sys.exit(0 if success else 1)
    else:
        print("用法: python qrcode_opencv.py <pdf_path> <output_path> [page_num]")
        print("示例: python qrcode_opencv.py license.pdf qrcode.png 0")
        sys.exit(1)