import json
import cv2
import os
import numpy as np

def crop_avatars():
    # Paths
    base_dir = r'd:\dsWork\aiData\static\Test'
    json_path = os.path.join(base_dir, 'data.json')
    images_dir = os.path.join(base_dir, 'Images')
    output_dir = os.path.join(base_dir, 'Result')

    # Create output directory if it doesn't exist
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"Created output directory: {output_dir}")

    # Load JSON data
    try:
        with open(json_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
    except Exception as e:
        print(f"Error loading JSON file: {e}")
        return

    print(f"Found {len(data)} entries in data.json")

    for item in data:
        filename = item['filename']
        x = int(item['x'])
        y = int(item['y'])
        r = int(item['r'])

        image_path = os.path.join(images_dir, filename)
        if not os.path.exists(image_path):
            print(f"Image not found: {image_path}")
            continue

        # Read image
        img = cv2.imread(image_path)
        if img is None:
            print(f"Failed to read image: {image_path}")
            continue

        height, width = img.shape[:2]

        # Calculate bounding box for cropping
        x1 = max(0, x - r)
        y1 = max(0, y - r)
        x2 = min(width, x + r)
        y2 = min(height, y + r)

        # Ensure valid crop dimensions
        if x2 <= x1 or y2 <= y1:
            print(f"Invalid crop dimensions for {filename}")
            continue

        # Create a mask for the circle
        # Mask needs to be the size of the original image first to handle edge cases correctly
        mask = np.zeros((height, width), dtype=np.uint8)
        cv2.circle(mask, (x, y), r, (255), -1)

        # Add alpha channel to the image
        b, g, r_channel = cv2.split(img)
        rgba = cv2.merge([b, g, r_channel, mask])

        # Crop the RGBA image
        cropped = rgba[y1:y2, x1:x2]

        # Since the crop might not be a perfect square if near edges, 
        # but the circle is centered at x,y.
        # Ideally we want the output to be a square bounding box of the circle.
        # If the circle goes out of bounds, the cropped image will be smaller.
        # Let's keep it simple: crop what is available. The alpha channel handles the shape.

        # Save result
        output_filename = os.path.splitext(filename)[0] + '.png'
        output_path = os.path.join(output_dir, output_filename)
        
        cv2.imwrite(output_path, cropped)
        print(f"Saved: {output_path}")

    print("Processing complete.")

if __name__ == "__main__":
    crop_avatars()