# -*- coding: utf-8 -*-
# Required libraries: paddleocr, paddlepaddle (or paddlepaddle-gpu)
# Install using: pip install paddlepaddle paddleocr
import os
import time
from paddleocr import PaddleOCR # Import PaddleOCR class
import logging

# Suppress excessive PaddleOCR logging if desired
logging.disable(logging.INFO) # Disables INFO level logs, keeps WARNING and ERROR
# logging.disable(logging.WARNING) # To disable WARNING level logs as well

# --- Configuration ---
BASE_IMAGE_DIR = '/Users/zpc01/workspace/zzlh/compliance/assets/images/dms1'
# Language configuration: 'ch' supports Chinese and English by default.
# Other options include 'en', 'korean', 'japan', 'french', 'german', etc.
LANG = 'ch'
USE_ANGLE_CLS = True # Use angle classification to help with rotated text
# PaddleOCR attempts to use GPU automatically if paddlepaddle-gpu is installed
# and CUDA is available. You can force CPU with use_gpu=False.
# USE_GPU = False # Uncomment to force CPU usage

# --- Initialize PaddleOCR Engine ---
# This needs to run only once to download and load models into memory.
print(f"正在初始化 PaddleOCR (语言: {LANG})... 这可能需要一些时间，特别是首次运行时。")
try:
    # Initialize PaddleOCR
    # If you want to explicitly control GPU usage, add use_gpu=USE_GPU parameter
    ocr_engine = PaddleOCR(use_angle_cls=USE_ANGLE_CLS, lang=LANG) # Default: use_gpu=True (auto-detect)
    print("PaddleOCR 初始化成功。")
except Exception as e:
    print(f"初始化 PaddleOCR 时出错: {e}")
    print("请确保已正确安装 paddleocr 和 paddlepaddle (或 paddlepaddle-gpu)。")
    print("GPU 用户请确保 CUDA/cuDNN 环境配置正确。")
    exit() # Exit if initialization fails

# --- Iterate through base directory's subfolders ---
print(f"开始处理目录: {BASE_IMAGE_DIR}")

if not os.path.isdir(BASE_IMAGE_DIR):
    print(f"错误：基础目录 '{BASE_IMAGE_DIR}' 不存在或不是一个有效的目录。")
    exit()

try:
    subfolders = [f for f in os.listdir(BASE_IMAGE_DIR) if os.path.isdir(os.path.join(BASE_IMAGE_DIR, f))]
except OSError as e:
    print(f"错误：无法访问目录 '{BASE_IMAGE_DIR}'。请检查权限。 {e}")
    exit()

if not subfolders:
    print(f"在 '{BASE_IMAGE_DIR}' 下没有找到任何子文件夹。")
    exit()

print(f"找到 {len(subfolders)} 个子文件夹，将逐一处理...")

# --- Process each subfolder ---
for folder_name in subfolders:
    subdir_path = os.path.join(BASE_IMAGE_DIR, folder_name)
    print(f"\n--- 正在处理文件夹: {folder_name} ---")

    all_texts_in_folder = [] # Store all recognized text from this folder
    png_filenames_sorted = [] # Store sorted PNG filenames

    # --- Find and sort PNG files in the current subdirectory ---
    try:
        all_items_in_subdir = os.listdir(subdir_path)
        # Filter for PNG files only and sort them
        png_filenames = [
            f for f in all_items_in_subdir
            if f.lower().endswith('.png') and os.path.isfile(os.path.join(subdir_path, f))
        ]
        png_filenames.sort() # Sort alphabetically (dictionary order)
        png_filenames_sorted = png_filenames

        if not png_filenames_sorted:
             print(f"  在文件夹 '{folder_name}' 中未找到 PNG 图片。")
             continue # Skip to the next folder

    except OSError as e:
        print(f"  错误：无法访问子文件夹 '{subdir_path}' 或读取其内容。跳过此文件夹。 {e}")
        continue # Skip to the next folder

    print(f"  找到 {len(png_filenames_sorted)} 个 PNG 文件，将按字典序处理...")

    # --- Process PNG files in sorted order ---
    for filename in png_filenames_sorted:
        image_path = os.path.join(subdir_path, filename)
        print(f"  正在识别图片 (顺序: {png_filenames_sorted.index(filename) + 1}/{len(png_filenames_sorted)}): {filename} ...")

        try:
            start_time = time.time()
            # Perform OCR using PaddleOCR
            # The result is typically a list where each item corresponds to a detected text line.
            # Format: [[[box], (text, confidence)], [[box], (text, confidence)], ...]
            # Sometimes it might be nested further, e.g., [page_result]
            # We usually need result[0] for a single image.
            result = ocr_engine.ocr(image_path, cls=USE_ANGLE_CLS)
            end_time = time.time()

            # Extract text from the result structure
            extracted_texts = []
            if result and result[0]: # Check if result is not None and the first element (page/image result) exists
                 for line in result[0]: # Iterate through detected lines in the first (and likely only) page/image
                    # line format is usually [[box_coords], (text, confidence)]
                    if line and len(line) == 2 and isinstance(line[1], tuple) and len(line[1]) >= 1:
                         extracted_texts.append(line[1][0]) # Get the text part

            if extracted_texts:
                all_texts_in_folder.extend(extracted_texts)
                print(f"    识别到 {len(extracted_texts)} 段文本，耗时: {end_time - start_time:.2f} 秒")
            else:
                print("    未识别到文本。")

        except Exception as e:
            # Catch potential errors during OCR processing for a single file
            print(f"  处理图片 '{filename}' 时发生错误: {e}")
            # Decide whether to stop or continue; here we continue
            continue

    # --- Write aggregated text to a file ---
    if all_texts_in_folder:
        output_txt_filename = f"{folder_name}.txt"
        output_txt_path = os.path.join(subdir_path, output_txt_filename)

        try:
            with open(output_txt_path, 'w', encoding='utf-8') as f:
                for text in all_texts_in_folder:
                    f.write(text + '\n') # Write each text segment on a new line
            print(f"  已将所有识别结果写入文件: {output_txt_path}")
        except IOError as e:
            print(f"  错误：无法写入输出文件 '{output_txt_path}'. {e}")
    elif png_filenames_sorted: # Found PNGs but extracted no text
         print(f"  文件夹 '{folder_name}' 中的 PNG 图片均未识别到文本，不创建 .txt 文件。")
    # If png_filenames_sorted is empty, the message was printed earlier.

print("\n--- 所有文件夹处理完毕 ---")