import os
import time
from PIL import Image

# 导入配置
from Config import target_width, image_quality, compressed_images_dir


def process_image(img_path, target_width=target_width):
    """
    处理图像
    :param img_path: 输入图像路径
    :param target_width: 目标宽度（默认587像素，根据用户提供的好用图片尺寸）
    :param output_path: 输出结果的文件路径（可选）
    :param raw_output_path: 原始OCR结果的输出文件路径（可选）
    :return: OCR识别结果
    """
    start_time = time.time()

    # 图片预处理 - 缩小图片尺寸到指定宽度
    img = Image.open(img_path)
    width, height = img.size
    print(f"原始图片尺寸: {width}x{height}")

    # 计算缩放比例，保持纵横比
    scale = target_width / width
    new_width = target_width
    new_height = int(height * scale)

    print(f"目标图片尺寸: {new_width}x{new_height}")

    # 缩放图像
    resized_img = img.resize((new_width, new_height), Image.LANCZOS)

    # 确保压缩图片目录存在
    os.makedirs(compressed_images_dir, exist_ok=True)
    
    # 获取原始文件名并创建新的输出路径
    file_name = os.path.basename(img_path)
    output_name = os.path.splitext(file_name)[0] + "_compressed.jpg"
    output_path = os.path.join(compressed_images_dir, output_name)
    
    # 保存缩小后的图片
    resized_img.convert('RGB').save(output_path, 'JPEG', quality=image_quality, optimize=True)
    print(f"图片已缩小: {width}x{height} -> {new_width}x{new_height}")

    # 释放内存
    del resized_img
    img.close()
    del img


if __name__ == "__main__":
    import glob
    from Config import images_dir
    
    # 处理Images目录下的所有PNG图片
    image_files = glob.glob(os.path.join(images_dir, "page_*.png"))
    
    print(f"找到{len(image_files)}个图片文件需要处理")
    
    for i, img_path in enumerate(image_files):
        print(f"[{i+1}/{len(image_files)}] 处理图片: {os.path.basename(img_path)}")
        process_image(img_path)
        
    print(f"所有图片处理完成，压缩后的图片保存在: {compressed_images_dir}")