#!/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
import shutil
import tempfile
import hashlib


def safe_path_for_opencv(file_path):
    """
    为OpenCV创建安全的文件路径（避免中文路径问题）
    
    参数:
        file_path: 原始文件路径
    
    返回:
        (temp_path, cleanup_func): 临时路径和清理函数
    """
    
    # 检查路径是否包含非ASCII字符
    try:
        file_path.encode('ascii')
        # 如果没有异常，说明路径是ASCII安全的
        return file_path, lambda: None
    except UnicodeEncodeError:
        # 路径包含非ASCII字符，需要创建临时副本
        pass
    
    if not os.path.exists(file_path):
        return file_path, lambda: None
    
    # 创建临时文件
    _, ext = os.path.splitext(file_path)
    
    # 使用文件内容的hash作为临时文件名，确保唯一性
    with open(file_path, 'rb') as f:
        file_hash = hashlib.md5(f.read()).hexdigest()[:8]
    
    temp_dir = tempfile.gettempdir()
    temp_filename = f"opencv_safe_{file_hash}{ext}"
    temp_path = os.path.join(temp_dir, temp_filename)
    
    # 复制文件到临时位置
    try:
        shutil.copy2(file_path, temp_path)
        print(f"创建临时文件: {temp_filename}")
        
        def cleanup():
            try:
                if os.path.exists(temp_path):
                    os.remove(temp_path)
                    print(f"清理临时文件: {temp_filename}")
            except:
                pass
        
        return temp_path, cleanup
    except Exception as e:
        print(f"创建临时文件失败: {e}")
        return file_path, lambda: None


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_safe(image_path, output_image_path='qrcode.png'):
    """
    从图像文件中提取二维码（处理中文路径）
    
    参数:
        image_path: 图像文件路径
        output_image_path: 输出二维码图片的路径
    
    返回:
        成功返回True和二维码数据,失败返回False和None
    """
    
    try:
        print(f"处理图像文件: {os.path.basename(image_path)}")
        
        # 创建OpenCV安全的路径
        safe_path, cleanup = safe_path_for_opencv(image_path)
        
        try:
            # 使用OpenCV读取图像
            img_cv = cv2.imread(safe_path)
            if img_cv is None:
                print(f"OpenCV无法读取图像，尝试使用PIL...")
                
                # 备用方案：使用PIL读取，然后转换为OpenCV格式
                try:
                    pil_img = Image.open(image_path)
                    img_cv = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
                    print("使用PIL成功读取图像")
                except Exception as e:
                    print(f"PIL也无法读取图像: {e}")
                    return False, None
            else:
                print("OpenCV成功读取图像")
            
            # 使用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]
                
                # 保存二维码图片（也需要处理输出路径）
                output_safe_path, output_cleanup = safe_path_for_opencv(output_image_path)
                
                try:
                    success = cv2.imwrite(output_safe_path, qr_image)
                    if success and output_safe_path != output_image_path:
                        # 如果使用了临时路径，复制回原位置
                        shutil.copy2(output_safe_path, output_image_path)
                        output_cleanup()
                    
                    print(f"成功提取二维码!")
                    print(f"二维码内容: {data}")
                    print(f"二维码图片已保存至: {output_image_path}")
                    
                    return True, data
                except Exception as e:
                    print(f"保存二维码图片失败: {e}")
                    return False, None
                finally:
                    output_cleanup()
            else:
                print("未在图像中检测到二维码")
                return False, None
                
        finally:
            cleanup()
        
    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
        return False, None


def extract_qr_from_pdf_safe(pdf_path, output_image_path='qrcode.png', page_num=0):
    """
    从PDF中提取二维码（处理中文路径）
    
    参数:
        pdf_path: PDF文件路径
        output_image_path: 输出二维码图片的路径
        page_num: 要处理的页码(从0开始),默认为第一页
    
    返回:
        成功返回True和二维码数据,失败返回False和None
    """
    
    try:
        print(f"处理PDF文件第{page_num + 1}页: {os.path.basename(pdf_path)}")
        
        # PyMuPDF通常能处理中文路径，但我们仍然做保护
        try:
            # 打开PDF文件
            pdf_document = fitz.open(pdf_path)
        except Exception as e:
            print(f"直接打开PDF失败，尝试临时文件方案: {e}")
            
            # 创建临时副本
            safe_path, cleanup = safe_path_for_opencv(pdf_path)
            try:
                pdf_document = fitz.open(safe_path)
            except Exception as e2:
                print(f"临时文件方案也失败: {e2}")
                cleanup()
                return False, None
        
        try:
            # 获取指定页面
            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]
                
                # 保存二维码图片
                output_safe_path, output_cleanup = safe_path_for_opencv(output_image_path)
                
                try:
                    success = cv2.imwrite(output_safe_path, qr_image)
                    if success and output_safe_path != output_image_path:
                        shutil.copy2(output_safe_path, output_image_path)
                        output_cleanup()
                    
                    print(f"成功提取二维码!")
                    print(f"二维码内容: {data}")
                    print(f"二维码图片已保存至: {output_image_path}")
                    
                    return True, data
                except Exception as e:
                    print(f"保存二维码图片失败: {e}")
                    return False, None
                finally:
                    output_cleanup()
            else:
                print("未在PDF中检测到二维码")
                return False, None
                
        finally:
            pdf_document.close()
        
    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
        return False, None


def extract_qr_universal_safe(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
    
    print(f"文件路径: {input_path}")
    print(f"文件大小: {os.path.getsize(input_path)} bytes")
    
    file_type = detect_file_type(input_path)
    
    if file_type == 'pdf':
        return extract_qr_from_pdf_safe(input_path, output_path, page_num)
    elif file_type == 'image':
        return extract_qr_from_image_safe(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_safe(input_path, output_path, page_num)
        
        # 返回状态码: 0=成功, 1=失败
        sys.exit(0 if success else 1)
    else:
        print("用法: python qrcode_unicode_safe.py <input_path> <output_path> [page_num]")
        print("支持格式: PDF, JPG, JPEG, PNG, BMP, TIFF")
        print("支持中文路径和文件名")
        print("示例: python qrcode_unicode_safe.py 营业执照.pdf qrcode.png 0")
        print("示例: python qrcode_unicode_safe.py 营业执照.jpg qrcode.png")
        sys.exit(1)