import os
import json
import argparse
import re
from PIL import Image

def analyze_directory(directory_path):
    """
    Analyzes a directory of PNG files to extract information about categories,
    image counts, and resolution, then saves it to a JSON file.
    """
    png_files_matched = []
    max_category_id = -1
    image_resolution = None
    
    # Ensure target_dir for JSON file is valid
    if not os.path.isdir(directory_path):
        print(f"Error: Input directory '{directory_path}' not found.")
        return

    print(f"Analyzing directory: {directory_path}")

    for filename in os.listdir(directory_path):
        if filename.lower().endswith(".png"):
            # Regex to match '数字_任意字符.png'
            match = re.match(r"(\d+)_(.*?)\.png", filename, re.IGNORECASE)
            if match:
                try:
                    category_id = int(match.group(1))
                    max_category_id = max(max_category_id, category_id)
                    png_files_matched.append(filename)

                    if image_resolution is None: # Try to get resolution from the first matched file
                        try:
                            with Image.open(os.path.join(directory_path, filename)) as img:
                                image_resolution = {"width": img.width, "height": img.height}
                        except Exception as e:
                            print(f"Warning: Could not read resolution from {filename}: {e}")
                except ValueError:
                    print(f"Warning: Could not parse category ID from filename {filename}.")
            # else:
            #     print(f"Info: File '{filename}' does not match the expected 'category_imageName.png' format.")

    if not png_files_matched:
        print("No PNG files matching the format 'category_imageName.png' found in the directory.")
        # Create an empty/default info.json if no files are found? Or just return?
        # For now, let's just return if no relevant files are found.
        return

    total_categories_val = max_category_id + 1 if max_category_id != -1 else 0
    
    images_per_category_val = 0.0
    if total_categories_val > 0:
        images_per_category_val = len(png_files_matched) / total_categories_val
    
    info_data = {
        "total_categories": total_categories_val,
        "images_per_category": round(images_per_category_val, 2), # Rounded to 2 decimal places
        "image_resolution": image_resolution if image_resolution else "N/A",
        "total_matched_png_files": len(png_files_matched)
    }

    json_file_path = os.path.join(directory_path, "info.json")
    try:
        with open(json_file_path, 'w') as f:
            json.dump(info_data, f, indent=4)
        print(f"Information successfully saved to {json_file_path}")
    except IOError as e:
        print(f"Error: Could not write JSON file to {json_file_path}: {e}")
    except Exception as e:
        print(f"An unexpected error occurred while writing JSON: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Analyzes a directory of PNG files (e.g., '0_0.png', '0_20.png') "
                    "and creates an info.json file with category counts, "
                    "images per category, and image resolution."
    )
    parser.add_argument(
        "directory",
        type=str,
        help="Path to the directory containing the PNG files."
    )
    args = parser.parse_args()

    analyze_directory(args.directory) 
