import os
from PIL import Image, ImageChops

def homogenize_gear_images(
    input_dir="./images/equipment",
    output_dir="./images/equipment_clean",
    target_size=(800, 500),      # Standardized card canvas size (WxH)
    padding_percent=0.12,         # 12% margin around the subject
    color_threshold=245,          # Treats pixels brighter than #F5F5F5 as background
    alpha_threshold=15            # Ignore transparent pixels
):
    """
    Trims background whitespace/transparency from drone photos, scales each subject 
    uniformly, and centers it on a clean white canvas.
    """
    os.makedirs(output_dir, exist_ok=True)
    valid_exts = ('.png', '.jpg', '.jpeg', '.webp', '.jfif')

    target_w, target_h = target_size
    max_w = int(target_w * (1 - 2 * padding_percent))
    max_h = int(target_h * (1 - 2 * padding_percent))

    files = [f for f in os.listdir(input_dir) if f.lower().endswith(valid_exts)]
    print(f"Found {len(files)} images in '{input_dir}'...")

    for fname in files:
        in_path = os.path.join(input_dir, fname)
        try:
            with Image.open(in_path) as img:
                rgba = img.convert("RGBA")
                r, g, b, a = rgba.split()

                # Build a mask identifying non-background pixels
                alpha_mask = a.point(lambda p: 255 if p > alpha_threshold else 0)
                r_mask = r.point(lambda p: 255 if p < color_threshold else 0)
                g_mask = g.point(lambda p: 255 if p < color_threshold else 0)
                b_mask = b.point(lambda p: 255 if p < color_threshold else 0)

                # Combine RGB & Alpha detection
                color_mask = ImageChops.add(r_mask, ImageChops.add(g_mask, b_mask))
                final_mask = ImageChops.multiply(color_mask, alpha_mask)

                bbox = final_mask.getbbox()

                if not bbox:
                    print(f"⚠️ Warning: Skipped '{fname}' (no subject detected or pure white image)")
                    continue

                # Step 1: Crop tightly to subject
                cropped = rgba.crop(bbox)
                crop_w, crop_h = cropped.size

                # Step 2: Scale subject proportionally to fit target dimensions with padding
                scale = min(max_w / crop_w, max_h / crop_h)
                new_w = max(1, int(crop_w * scale))
                new_h = max(1, int(crop_h * scale))

                resized = cropped.resize((new_w, new_h), Image.Resampling.LANCZOS)

                # Step 3: Paste onto a clean, solid white canvas
                canvas = Image.new("RGBA", target_size, (255, 255, 255, 255))
                offset_x = (target_w - new_w) // 2
                offset_y = (target_h - new_h) // 2
                canvas.paste(resized, (offset_x, offset_y), resized)

                # Step 4: Export standardized JPG
                base_name = os.path.splitext(fname)[0]
                out_path = os.path.join(output_dir, f"{base_name}.jpg")
                canvas.convert("RGB").save(out_path, "JPEG", quality=92)

                print(f"✅ Processed: {fname} -> {out_path}")

        except Exception as e:
            print(f"❌ Error processing '{fname}': {e}")

if __name__ == "__main__":
    homogenize_gear_images()