import os
import sys
from PIL import Image

# --- Configuration ---
ALLOWED_EXTENSIONS = ['.png'] # Only PNG supports the transparency needed for this fix
# ---------------------

def extend_edge_color(input_path):
    """
    Finds fully transparent pixels bordering an opaque pixel and fills the 
    transparent pixel with the neighbor's color, making the new pixel opaque.
    """
    
    print(f"Processing image: {input_path}")
    
    try:
        # Open the image and ensure it has an alpha channel (RGBA)
        img = Image.open(input_path).convert("RGBA")
    except FileNotFoundError:
        return False, "Error: File not found."
    except Exception as e:
        return False, f"Error opening image: {e}"

    width, height = img.size
    
    # Get pixels list for modification
    data = list(img.getdata())
    
    # Create a list to store modifications; we don't modify 'data' while iterating
    # We will store (x, y, new_color_rgb) tuples
    modifications = []
    
    # --- Pass 1: Identify and store modifications ---
    
    # Iterate over every pixel
    for y in range(height):
        for x in range(width):
            
            # Index in the flattened data list
            index = y * width + x
            
            # The current pixel's alpha value
            current_alpha = data[index][3]
            
            # We ONLY care about pixels that are FULLY TRANSPARENT
            if current_alpha == 0:
                
                # Check neighbors (8 surrounding pixels)
                
                # Define neighbors offsets (dx, dy) including self-correction for image edges
                neighbors = [
                    (-1, 0), (1, 0), (0, -1), (0, 1), # Cardinal
                    (-1, -1), (1, -1), (-1, 1), (1, 1) # Diagonal
                ]
                
                found_opaque_neighbor = False
                
                for dx, dy in neighbors:
                    nx, ny = x + dx, y + dy
                    
                    # Ensure neighbor coordinates are within the image bounds
                    if 0 <= nx < width and 0 <= ny < height:
                        neighbor_index = ny * width + nx
                        neighbor_pixel = data[neighbor_index]
                        neighbor_alpha = neighbor_pixel[3]
                        
                        # If the neighbor is FULLY OPAQUE (Alpha = 255)
                        if neighbor_alpha == 255:
                            # Use the neighbor's RGB color
                            new_color_rgb = neighbor_pixel[:3] 
                            
                            # Store the location and new color to apply later
                            modifications.append((index, new_color_rgb))
                            found_opaque_neighbor = True
                            break # Move to the next transparent pixel
                
    # --- Pass 2: Apply modifications ---
    processed_count = 0
    for index, new_color_rgb in modifications:
        # Create the new opaque pixel: (R, G, B, 255)
        data[index] = new_color_rgb + (255,)
        processed_count += 1

    if processed_count == 0:
        return False, "Skipped: No fully transparent pixels were found bordering an opaque area."
        
    # Put the modified data back into the image
    img.putdata(data)

    # 3. Save the image, OVERWRITING THE ORIGINAL FILE
    try:
        img.save(input_path)
        return True, f"OVERWRITTEN: Extended edge color for {processed_count} transparent pixels."
    except Exception as e:
        return False, f"Error saving image: {e}"


def batch_edge_color_interactive():
    """Interactively finds and extends edge colors in images in subdirectories, overwriting originals."""
    
    # --- SAFETY CHECK ---
    print("\n!!! WARNING: THIS SCRIPT WILL OVERWRITE YOUR ORIGINAL FILES AND REMOVE TRANSPARENT BORDERS !!!")
    input("Press Enter to continue or Ctrl+C to stop and make a backup...")
    # --------------------
    
    start_directory = os.getcwd()
    total_processed_count = 0
    
    for root, dirs, files in os.walk(start_directory):
        
        # Skip the starting directory—it must always be processed if we start the script.
        if root != start_directory:
            
            # --- Interactive Prompt ---
            relative_path = os.path.relpath(root, start_directory)
            
            has_images = any(f.lower().endswith(tuple(ALLOWED_EXTENSIONS)) for f in files)
            
            if not has_images:
                continue

            response = input(f"\nExtend edge color in directory '{relative_path}'? (y/n): ").lower()
            
            if response != 'y':
                print(f"Skipping directory: {relative_path}")
                continue
            # --- End Interactive Prompt ---
        
        
        # Process files within the current root (folder)
        processed_in_folder = 0
        
        for filename in files:
            name, ext = os.path.splitext(filename)
            
            # Check if the file is an allowed image and not the script itself
            if ext.lower() in ALLOWED_EXTENSIONS and filename != os.path.basename(sys.argv[0]):
                
                input_path = os.path.join(root, filename)
                
                print(f"  Processing: {filename}...")

                # Call the core fill function
                success, message = extend_edge_color(input_path)
                
                # Report the result
                if success:
                    print(f"  ✅ {message}.")
                    total_processed_count += 1
                    processed_in_folder += 1
                else:
                    print(f"  ❌ Failed: {message}")

        if processed_in_folder > 0:
            print(f"--- Folder Complete: {processed_in_folder} files processed in {os.path.relpath(root, start_directory)} ---")


    print(f"\n*** BATCH EDGE EXTENSION FINISHED ***")
    print(f"Total images processed: {total_processed_count}")


if __name__ == "__main__":
    try:
        # Check Pillow version compatibility
        _ = Image.NEAREST 
        batch_edge_color_interactive()
    except AttributeError:
        print("\nFATAL ERROR: Your Pillow library version is too old. Please update it ('pip install --upgrade Pillow').")
    except ImportError:
        print("\nFATAL ERROR: The Pillow library is not installed. Please install it ('pip install Pillow').")
    except Exception as e:
        print(f"\nAn unexpected runtime error occurred: {e}")