#!/usr/bin/env python3
# -*- coding: utf-8 -*-

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


def detect_file_type(file_path):
    """检测文件类型"""
    _, ext = os.path.splitext(file_path.lower())
    
    if ext in ['.pdf']:
        return 'pdf'
    elif ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']:
        return 'image'
    else:
        return 'unknown'


def extract_qr_from_image_opencv(image_path, output_image_path='qrcode.png'):
    """
    从图像文件中提取二维码（使用OpenCV）
    
    参数:
        image_path: 图像文件路径
        output_image_path: 输出二维码图片的路径
    
    返回:
        成功返回True和二维码数据,失败返回False和None
    """
    
    try:
        # 直接读取图像
        img_cv = cv2.imread(image_path)
        if img_cv is None:
            print(f"无法读取图像文件: {image_path}")
            return False, None
        
        # 使用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}")
            
            return True, data
        else:
            print("未在图像中检测到二维码")
            return False, None
        
    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
        return False, None


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_qr_universal(input_path, output_path='qrcode.png', page_num=0):
    """
    通用二维码提取函数，自动检测文件类型
    
    参数:
        input_path: 输入文件路径（PDF或图像）
        output_path: 输出二维码图片的路径
        page_num: PDF页码（仅对PDF有效）
    
    返回:
        成功返回True和二维码数据,失败返回False和None
    """
    
    if not os.path.exists(input_path):
        print(f"输入文件不存在: {input_path}")
        return False, None
    
    file_type = detect_file_type(input_path)
    
    if file_type == 'pdf':
        print(f"检测到PDF文件，处理第{page_num + 1}页...")
        return extract_qr_from_pdf_opencv(input_path, output_path, page_num)
    elif file_type == 'image':
        print(f"检测到图像文件，直接处理...")
        return extract_qr_from_image_opencv(input_path, output_path)
    else:
        print(f"不支持的文件类型: {input_path}")
        return False, None


# 使用示例
if __name__ == "__main__":
    # 支持命令行调用
    if len(sys.argv) >= 3:
        input_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_universal(input_path, output_path, page_num)
        
        # 返回状态码: 0=成功, 1=失败
        sys.exit(0 if success else 1)
    else:
        print("用法: python qrcode_universal.py <input_path> <output_path> [page_num]")
        print("支持格式: PDF, JPG, JPEG, PNG, BMP, TIFF")
        print("示例: python qrcode_universal.py license.pdf qrcode.png 0")
        print("示例: python qrcode_universal.py license.jpg qrcode.png")
        sys.exit(1)