import cv2
import os
import glob
import numpy as np

def apply_mask_and_crop(input_folder, mask, output_folder):
    """
    Apply binary mask to all images, crop to masked region, and save to output folder
    
    Args:
        input_folder (str): Path to folder containing input images
        mask_path (str): Path to binary mask image
        output_folder (str): Path to save cropped masked images
    """
    # Load and prepare mask
    # mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    
    if mask is None:
        raise ValueError(f"Could not load mask from {mask_path}")
    
    _, binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
    
    # Create output directory if it doesn't exist
    os.makedirs(output_folder, exist_ok=True)
    
    # Get list of image files
    image_files = glob.glob(os.path.join(input_folder, "*.jpg")) + \
                 glob.glob(os.path.join(input_folder, "*.png")) + \
                 glob.glob(os.path.join(input_folder, "*.bmp"))
    
    if not image_files:
        print(f"No images found in {input_folder}")
        return
    
    print(f"Found {len(image_files)} images to process")
    i = 0
    for img_path in image_files:
        # Load image
        img = cv2.imread(img_path)
        if img is None:
            print(f"Warning: Could not read image {img_path}")
            continue
        
        # Resize mask if dimensions don't match
        if img.shape[:2] != binary_mask.shape[:2]:
            resized_mask = cv2.resize(binary_mask, (img.shape[1], img.shape[0]))
        else:
            resized_mask = binary_mask
        
        # Apply mask
        masked_img = cv2.bitwise_and(img, img, mask=resized_mask)
        
        # Find contours to get bounding box of mask
        contours, _ = cv2.findContours(resized_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if not contours:
            print(f"No mask area found in {img_path} - skipping")
            continue
        
        # Get bounding rectangle of largest contour
        x, y, w, h = cv2.boundingRect(max(contours, key=cv2.contourArea))
        
        # Crop to masked region
        cropped_img = masked_img[y:y+h, x:x+w]
        
        # Create output path (preserve original filename)
        filename = os.path.basename(img_path)
        output_path = os.path.join(output_folder, filename)
        
        # Save cropped image
        cv2.imwrite(output_path, cropped_img)
        if i == 0:
            preview_path = os.path.join(output_folder, f"preview_{filename}")
            cv2.imwrite(preview_path, cropped_img)
            return preview_path
        i= i+1
        # print(f"Processed and saved: {output_path}")
    
    print(f"\nProcessing complete! Saved {len(image_files)} cropped images to {output_folder}")