"""
Gemini Watermark Removal Tool
Removes visible Gemini watermarks using reverse alpha blending.

Based on the mathematical formula:
    watermarked = α × logo + (1 - α) × original
    
Solving for original:
    original = (watermarked - α × logo) / (1 - α)
Python port with embedded alpha maps.

Usage:
    python GWMRTool.py image.jpg
    python GWMRTool.py -i input.jpg -o output.jpg
    python GWMRTool.py -i ./input_folder/ -o ./output_folder/
"""

import numpy as np
from PIL import Image
import argparse
from pathlib import Path
import sys
import os
from typing import Tuple, Optional


# ============================================================================
# EMBEDDED ALPHA MAPS
# ============================================================================

# Default paths for alpha maps (relative to script location)
SCRIPT_DIR = Path(__file__).parent.resolve()
DEFAULT_ALPHA_MAP_SMALL_PATH = SCRIPT_DIR / "bg_48.png"
DEFAULT_ALPHA_MAP_LARGE_PATH = SCRIPT_DIR / "bg_96.png"


class AlphaMapManager:
    """
    Manages loading and caching of alpha maps.
    """
    
    _cache = {}
    
    @classmethod
    def load_alpha_map(cls, path: Path, expected_size: int) -> np.ndarray:
        """
        Load alpha map from file with caching.
        
        Args:
            path: Path to alpha map image
            expected_size: Expected size (48 or 96)
            
        Returns:
            Alpha map as float32 array [0, 1]
        """
        cache_key = str(path)
        
        if cache_key in cls._cache:
            return cls._cache[cache_key]
        
        if not path.exists():
            raise FileNotFoundError(f"Alpha map not found: {path}")
        
        # Load image
        img = Image.open(path)
        img_array = np.array(img)
        
        # Calculate alpha map using max of RGB channels
        if len(img_array.shape) == 3:
            if img_array.shape[2] >= 3:
                # Use max of RGB channels for brightness
                gray = np.max(img_array[:, :, :3], axis=2)
            else:
                gray = img_array[:, :, 0]
        else:
            gray = img_array
        
        # Normalize to [0, 1]
        alpha_map = gray.astype(np.float32) / 255.0
        
        # Validate size
        if alpha_map.shape[0] != expected_size or alpha_map.shape[1] != expected_size:
            print(f"Warning: Alpha map size {alpha_map.shape} doesn't match expected {expected_size}x{expected_size}")
            # Resize if needed
            img_resized = Image.fromarray(gray).resize(
                (expected_size, expected_size), 
                Image.Resampling.LANCZOS
            )
            alpha_map = np.array(img_resized).astype(np.float32) / 255.0
        
        # Cache and return
        cls._cache[cache_key] = alpha_map
        return alpha_map
    
    @classmethod
    def clear_cache(cls):
        """Clear the alpha map cache."""
        cls._cache.clear()


# ============================================================================
# BLEND MODES (Translated from C++ blend_modes.cpp)
# ============================================================================

class BlendModes:
    """
    Alpha blending operations for watermark removal/addition.
    Direct translation from C++ blend_modes.cpp
    """
    
    # Constants from C++ code
    ALPHA_THRESHOLD = 0.002  # Ignore very small alpha (noise)
    MAX_ALPHA = 0.99         # Avoid division by near-zero
    
    @staticmethod
    def calculate_alpha_map(bg_capture: np.ndarray) -> np.ndarray:
        """
        Calculate alpha map from background capture.
        Uses max of RGB channels for brightness (handles colored logos).
        
        Args:
            bg_capture: Background capture image (grayscale or RGB/RGBA)
            
        Returns:
            Alpha map as float32 array normalized to [0, 1]
        """
        if bg_capture is None or bg_capture.size == 0:
            raise ValueError("Empty background capture")
        
        # Handle different channel configurations
        if len(bg_capture.shape) == 3:
            if bg_capture.shape[2] >= 3:
                # RGB/RGBA: Use max of RGB channels for brightness
                gray = np.max(bg_capture[:, :, :3], axis=2)
            else:
                gray = bg_capture[:, :, 0]
        else:
            # Already grayscale
            gray = bg_capture.copy()
        
        # Convert to float and normalize to [0, 1]
        alpha_map = gray.astype(np.float32) / 255.0
        
        return alpha_map
    
    @staticmethod
    def remove_watermark_alpha_blend(
        image: np.ndarray,
        alpha_map: np.ndarray,
        position: Tuple[int, int],
        logo_value: float = 255.0
    ) -> np.ndarray:
        """
        Remove watermark using reverse alpha blending (vectorized).
        
        Gemini applies: watermarked = alpha * logo + (1 - alpha) * original
        We reverse:     original = (watermarked - alpha * logo) / (1 - alpha)
        
        Args:
            image: Input image (RGB, uint8)
            alpha_map: Alpha map (float32, normalized to [0, 1])
            position: (x, y) position of watermark top-left corner
            logo_value: Logo pixel value (default 255 for white logo)
            
        Returns:
            Image with watermark removed
        """
        if image is None or image.size == 0:
            raise ValueError("Empty image")
        if alpha_map is None or alpha_map.size == 0:
            raise ValueError("Empty alpha map")
        if len(image.shape) != 3 or image.shape[2] != 3:
            raise ValueError("Image must be RGB (3 channels)")
        if alpha_map.dtype != np.float32:
            alpha_map = alpha_map.astype(np.float32)
        
        # Make a copy to avoid modifying original
        result = image.copy()
        
        # Calculate region of interest
        x, y = position
        h, w = alpha_map.shape[:2]
        
        # Clip to image bounds
        x1 = max(0, x)
        y1 = max(0, y)
        x2 = min(image.shape[1], x + w)
        y2 = min(image.shape[0], y + h)
        
        if x1 >= x2 or y1 >= y2:
            return result
        
        # Calculate ROI offsets
        alpha_x1 = x1 - x
        alpha_y1 = y1 - y
        alpha_x2 = alpha_x1 + (x2 - x1)
        alpha_y2 = alpha_y1 + (y2 - y1)
        
        # Get ROIs
        alpha_region = alpha_map[alpha_y1:alpha_y2, alpha_x1:alpha_x2]
        image_region = result[y1:y2, x1:x2]
        
        # Convert image region to float
        image_f = image_region.astype(np.float32)
        
        # Create mask for pixels to process (alpha >= threshold)
        process_mask = alpha_region >= BlendModes.ALPHA_THRESHOLD
        
        # Clamp alpha values
        alpha_clamped = np.minimum(alpha_region, BlendModes.MAX_ALPHA)
        one_minus_alpha = 1.0 - alpha_clamped
        
        # Avoid division by zero
        one_minus_alpha = np.maximum(one_minus_alpha, 1e-6)
        
        # Expand alpha to 3 channels for broadcasting
        alpha_3c = alpha_clamped[:, :, np.newaxis]
        one_minus_alpha_3c = one_minus_alpha[:, :, np.newaxis]
        process_mask_3c = process_mask[:, :, np.newaxis]
        
        # Reverse alpha blending formula:
        # original = (watermarked - alpha * logo) / (1 - alpha)
        original = (image_f - alpha_3c * logo_value) / one_minus_alpha_3c
        original = np.clip(original, 0.0, 255.0)
        
        # Only apply to pixels that need processing
        image_f = np.where(process_mask_3c, original, image_f)
        
        # Convert back to 8-bit
        result[y1:y2, x1:x2] = image_f.astype(np.uint8)
        
        return result
    
    @staticmethod
    def add_watermark_alpha_blend(
        image: np.ndarray,
        alpha_map: np.ndarray,
        position: Tuple[int, int],
        logo_value: float = 255.0
    ) -> np.ndarray:
        """
        Add watermark using alpha blending (same as Gemini).
        
        Formula: result = alpha * logo + (1 - alpha) * original
        
        Args:
            image: Input image (RGB, uint8)
            alpha_map: Alpha map (float32, normalized to [0, 1])
            position: (x, y) position of watermark top-left corner
            logo_value: Logo pixel value (default 255 for white logo)
            
        Returns:
            Image with watermark added
        """
        if image is None or image.size == 0:
            raise ValueError("Empty image")
        if alpha_map is None or alpha_map.size == 0:
            raise ValueError("Empty alpha map")
        if len(image.shape) != 3 or image.shape[2] != 3:
            raise ValueError("Image must be RGB (3 channels)")
        if alpha_map.dtype != np.float32:
            alpha_map = alpha_map.astype(np.float32)
        
        # Make a copy to avoid modifying original
        result = image.copy()
        
        # Calculate region of interest
        x, y = position
        h, w = alpha_map.shape[:2]
        
        # Clip to image bounds
        x1 = max(0, x)
        y1 = max(0, y)
        x2 = min(image.shape[1], x + w)
        y2 = min(image.shape[0], y + h)
        
        if x1 >= x2 or y1 >= y2:
            return result
        
        # Calculate ROI offsets
        alpha_x1 = x1 - x
        alpha_y1 = y1 - y
        alpha_x2 = alpha_x1 + (x2 - x1)
        alpha_y2 = alpha_y1 + (y2 - y1)
        
        # Get ROIs
        alpha_region = alpha_map[alpha_y1:alpha_y2, alpha_x1:alpha_x2]
        image_region = result[y1:y2, x1:x2]
        
        # Convert image region to float
        image_f = image_region.astype(np.float32)
        
        # Create mask for pixels to process
        process_mask = alpha_region >= BlendModes.ALPHA_THRESHOLD
        
        # Expand alpha to 3 channels for broadcasting
        alpha_3c = alpha_region[:, :, np.newaxis]
        one_minus_alpha_3c = (1.0 - alpha_region)[:, :, np.newaxis]
        process_mask_3c = process_mask[:, :, np.newaxis]
        
        # Alpha blending formula:
        # result = alpha * logo + (1 - alpha) * original
        blended = alpha_3c * logo_value + one_minus_alpha_3c * image_f
        blended = np.clip(blended, 0.0, 255.0)
        
        # Only apply to pixels that need processing
        image_f = np.where(process_mask_3c, blended, image_f)
        
        # Convert back to 8-bit
        result[y1:y2, x1:x2] = image_f.astype(np.uint8)
        
        return result


# ============================================================================
# WATERMARK REMOVER
# ============================================================================

class GeminiWatermarkRemover:
    """
    Removes Gemini watermarks using reverse alpha blending.
    Uses bg_48.png and bg_96.png alpha maps by default.
    """
    
    # Watermark size configurations
    SMALL_SIZE = 48
    LARGE_SIZE = 96
    SMALL_MARGIN = 32
    LARGE_MARGIN = 64
    SIZE_THRESHOLD = 1024
    
    # Default logo value (white)
    DEFAULT_LOGO_VALUE = 255.0
    
    def __init__(self, 
                 alpha_map_small_path: Path = None,
                 alpha_map_large_path: Path = None,
                 logo_value: float = DEFAULT_LOGO_VALUE):
        """
        Initialize the watermark remover.
        
        Args:
            alpha_map_small_path: Path to 48x48 alpha map (default: bg_48.png)
            alpha_map_large_path: Path to 96x96 alpha map (default: bg_96.png)
            logo_value: Logo pixel value (default 255 for white)
        """
        # Use default paths if not specified
        self.alpha_map_small_path = alpha_map_small_path or DEFAULT_ALPHA_MAP_SMALL_PATH
        self.alpha_map_large_path = alpha_map_large_path or DEFAULT_ALPHA_MAP_LARGE_PATH
        self.logo_value = logo_value
        
        # Lazy-loaded alpha maps
        self._alpha_map_small = None
        self._alpha_map_large = None
    
    @property
    def alpha_map_small(self) -> np.ndarray:
        """Lazy-load small alpha map."""
        if self._alpha_map_small is None:
            self._alpha_map_small = AlphaMapManager.load_alpha_map(
                Path(self.alpha_map_small_path), 
                self.SMALL_SIZE
            )
        return self._alpha_map_small
    
    @property
    def alpha_map_large(self) -> np.ndarray:
        """Lazy-load large alpha map."""
        if self._alpha_map_large is None:
            self._alpha_map_large = AlphaMapManager.load_alpha_map(
                Path(self.alpha_map_large_path), 
                self.LARGE_SIZE
            )
        return self._alpha_map_large
    
    def detect_watermark_size(self, image: np.ndarray) -> Tuple[int, int]:
        """
        Auto-detect watermark size based on image dimensions.
        
        Args:
            image: Input image as numpy array
            
        Returns:
            Tuple of (size, margin) for watermark
        """
        height, width = image.shape[:2]
        
        if width <= self.SIZE_THRESHOLD or height <= self.SIZE_THRESHOLD:
            return self.SMALL_SIZE, self.SMALL_MARGIN
        else:
            return self.LARGE_SIZE, self.LARGE_MARGIN
    
    def get_watermark_position(self, image: np.ndarray, 
                                size: int, margin: int) -> Tuple[int, int]:
        """
        Get the (x, y) position of the watermark top-left corner.
        Watermark is in bottom-right corner.
        
        Args:
            image: Input image
            size: Watermark size (48 or 96)
            margin: Margin from edge
            
        Returns:
            Tuple of (x, y) position
        """
        height, width = image.shape[:2]
        
        x = width - margin - size
        y = height - margin - size
        
        return x, y
    
    def remove_watermark(self, image: np.ndarray, 
                         force_size: str = None,
                         verbose: bool = False) -> np.ndarray:
        """
        Remove watermark from image.
        
        Args:
            image: Input image as numpy array (RGB)
            force_size: Force 'small' (48x48) or 'large' (96x96)
            verbose: Print debug information
            
        Returns:
            Image with watermark removed
        """
        # Ensure RGB
        if len(image.shape) == 2:
            # Grayscale - convert to RGB
            image = np.stack([image] * 3, axis=-1)
        elif len(image.shape) == 3 and image.shape[2] == 4:
            # RGBA - drop alpha
            image = image[:, :, :3]
        elif len(image.shape) != 3 or image.shape[2] != 3:
            raise ValueError("Image must be RGB or RGBA")
        
        # Detect or force watermark size
        if force_size == 'small':
            size, margin = self.SMALL_SIZE, self.SMALL_MARGIN
        elif force_size == 'large':
            size, margin = self.LARGE_SIZE, self.LARGE_MARGIN
        else:
            size, margin = self.detect_watermark_size(image)
        
        if verbose:
            print(f"  Watermark size: {size}x{size}, margin: {margin}px")
        
        # Get watermark position
        x, y = self.get_watermark_position(image, size, margin)
        
        if verbose:
            print(f"  Watermark position: ({x}, {y})")
        
        # Validate position
        if x < 0 or y < 0:
            if verbose:
                print(f"  Warning: Image too small for {size}x{size} watermark")
            return image.copy()
        
        # Get appropriate alpha map
        if size == self.SMALL_SIZE:
            alpha_map = self.alpha_map_small
        else:
            alpha_map = self.alpha_map_large
        
        # Remove watermark
        result = BlendModes.remove_watermark_alpha_blend(
            image=image,
            alpha_map=alpha_map,
            position=(x, y),
            logo_value=self.logo_value
        )
        
        return result
    
    def add_watermark(self, image: np.ndarray,
                      force_size: str = None,
                      verbose: bool = False) -> np.ndarray:
        """
        Add watermark to image (for testing purposes).
        
        Args:
            image: Input image as numpy array (RGB)
            force_size: Force 'small' (48x48) or 'large' (96x96)
            verbose: Print debug information
            
        Returns:
            Image with watermark added
        """
        # Ensure RGB
        if len(image.shape) == 2:
            image = np.stack([image] * 3, axis=-1)
        elif len(image.shape) == 3 and image.shape[2] == 4:
            image = image[:, :, :3]
        elif len(image.shape) != 3 or image.shape[2] != 3:
            raise ValueError("Image must be RGB or RGBA")
        
        # Detect or force watermark size
        if force_size == 'small':
            size, margin = self.SMALL_SIZE, self.SMALL_MARGIN
        elif force_size == 'large':
            size, margin = self.LARGE_SIZE, self.LARGE_MARGIN
        else:
            size, margin = self.detect_watermark_size(image)
        
        if verbose:
            print(f"  Watermark size: {size}x{size}, margin: {margin}px")
        
        # Get watermark position
        x, y = self.get_watermark_position(image, size, margin)
        
        # Validate position
        if x < 0 or y < 0:
            if verbose:
                print(f"  Warning: Image too small for {size}x{size} watermark")
            return image.copy()
        
        # Get appropriate alpha map
        if size == self.SMALL_SIZE:
            alpha_map = self.alpha_map_small
        else:
            alpha_map = self.alpha_map_large
        
        # Add watermark
        result = BlendModes.add_watermark_alpha_blend(
            image=image,
            alpha_map=alpha_map,
            position=(x, y),
            logo_value=self.logo_value
        )
        
        return result


# ============================================================================
# FILE OPERATIONS
# ============================================================================

def load_image(path: str) -> np.ndarray:
    """Load image as RGB numpy array."""
    img = Image.open(path).convert('RGB')
    return np.array(img)


def save_image(image: np.ndarray, path: str, quality: int = 95):
    """Save numpy array as image."""
    img = Image.fromarray(image)
    
    # Determine format from extension
    ext = Path(path).suffix.lower()
    
    if ext in ['.jpg', '.jpeg']:
        img.save(path, 'JPEG', quality=quality)
    elif ext == '.png':
        img.save(path, 'PNG')
    elif ext == '.webp':
        img.save(path, 'WEBP', quality=quality)
    elif ext == '.bmp':
        img.save(path, 'BMP')
    else:
        img.save(path)


def process_single_file(input_path: str, 
                        output_path: str = None,
                        remover: GeminiWatermarkRemover = None,
                        force_size: str = None,
                        verbose: bool = False,
                        add_mode: bool = False):
    """
    Process a single image file.
    
    Args:
        input_path: Path to input image
        output_path: Path for output (None = overwrite input)
        remover: Watermark remover instance
        force_size: Force 'small' or 'large'
        verbose: Print verbose output
        add_mode: Add watermark instead of removing
    """
    if output_path is None:
        output_path = input_path
    
    if verbose:
        print(f"Processing: {input_path}")
    
    # Load image
    image = load_image(input_path)
    
    if verbose:
        print(f"  Image size: {image.shape[1]}x{image.shape[0]}")
    
    # Process
    if add_mode:
        result = remover.add_watermark(image, force_size, verbose)
        action = "Added watermark"
    else:
        result = remover.remove_watermark(image, force_size, verbose)
        action = "Removed watermark"
    
    # Save result
    save_image(result, output_path)
    
    if verbose:
        print(f"  {action}, saved to: {output_path}")


def process_directory(input_dir: str, 
                      output_dir: str,
                      remover: GeminiWatermarkRemover = None,
                      force_size: str = None,
                      verbose: bool = False,
                      add_mode: bool = False):
    """
    Process all images in a directory.
    
    Args:
        input_dir: Input directory path
        output_dir: Output directory path
        remover: Watermark remover instance
        force_size: Force 'small' or 'large'
        verbose: Print verbose output
        add_mode: Add watermark instead of removing
    """
    input_path = Path(input_dir)
    output_path = Path(output_dir)
    
    # Create output directory
    output_path.mkdir(parents=True, exist_ok=True)
    
    # Supported formats
    formats = {'.jpg', '.jpeg', '.png', '.webp', '.bmp'}
    
    # Find all images
    images = [f for f in input_path.iterdir() 
              if f.is_file() and f.suffix.lower() in formats]
    
    total = len(images)
    if verbose:
        print(f"Found {total} images to process")
    
    processed = 0
    errors = 0
    
    for idx, img_path in enumerate(images, 1):
        out_file = output_path / img_path.name
        try:
            if verbose:
                print(f"\n[{idx}/{total}] ", end="")
            process_single_file(
                str(img_path), str(out_file), 
                remover, force_size, verbose, add_mode
            )
            processed += 1
        except Exception as e:
            print(f"Error processing {img_path.name}: {e}")
            errors += 1
    
    print(f"\nCompleted: {processed} processed, {errors} errors")


# ============================================================================
# CLI
# ============================================================================

def print_banner():
    """Print ASCII banner."""
    banner = r"""
╔═══════════════════════════════════════════════════════════════════╗
║                                                                   ║
║         ██████╗ ███████╗███╗   ███╗██╗███╗   ██╗██╗               ║
║        ██╔════╝ ██╔════╝████╗ ████║██║████╗  ██║██║               ║
║        ██║  ███╗█████╗  ██╔████╔██║██║██╔██╗ ██║██║               ║
║        ██║   ██║██╔══╝  ██║╚██╔╝██║██║██║╚██╗██║██║               ║
║        ╚██████╔╝███████╗██║ ╚═╝ ██║██║██║ ╚████║██║               ║
║         ╚═════╝ ╚══════╝╚═╝     ╚═╝╚═╝╚═╝  ╚═══╝╚═╝               ║
║                                                                   ║
║            Watermark Removal Tool (Python Edition)                ║
║              Reverse Alpha Blending Technology                    ║
║                                                                   ║
╚═══════════════════════════════════════════════════════════════════╝
    """
    print(banner)


def check_alpha_maps():
    """Check if alpha map files exist and print status."""
    small_exists = DEFAULT_ALPHA_MAP_SMALL_PATH.exists()
    large_exists = DEFAULT_ALPHA_MAP_LARGE_PATH.exists()
    
    print("Alpha Map Status:")
    print(f"  bg_48.png (48x48): {'✓ Found' if small_exists else '✗ Not found'} - {DEFAULT_ALPHA_MAP_SMALL_PATH}")
    print(f"  bg_96.png (96x96): {'✓ Found' if large_exists else '✗ Not found'} - {DEFAULT_ALPHA_MAP_LARGE_PATH}")
    print()
    
    if not small_exists or not large_exists:
        print("WARNING: Missing alpha map files!")
        print("Please ensure bg_48.png and bg_96.png are in the same directory as this script.")
        print()
        return False
    
    return True


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description='Remove Gemini watermarks using reverse alpha blending',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
Examples:
  # Simple mode - edit in place
  python GWMRTool.py image.jpg
  
  # Specify output file
  python GWMRTool.py -i input.jpg -o output.jpg
  
  # Batch processing
  python GWMRTool.py -i ./input_folder/ -o ./output_folder/
  
  # Force specific watermark size
  python GWMRTool.py -i image.jpg -o output.jpg --force-large
  
  # Add watermark (for testing)
  python GWMRTool.py -i clean.jpg -o watermarked.jpg --add

Required Files:
  bg_48.png  - 48x48 alpha map (for images <= 1024px)
  bg_96.png  - 96x96 alpha map (for images > 1024px)
  
  Place these files in the same directory as this script.

Mathematical Formula:
  Watermark Application:  watermarked = α × logo + (1 - α) × original
  Reverse (Removal):      original = (watermarked - α × logo) / (1 - α)
'''
    )
    
    parser.add_argument('simple_input', nargs='?', 
                        help='Image file (simple mode - edits in place)')
    parser.add_argument('-i', '--input', 
                        help='Input image file or directory')
    parser.add_argument('-o', '--output', 
                        help='Output image file or directory')
    parser.add_argument('--force-small', action='store_true',
                        help='Force 48x48 watermark size')
    parser.add_argument('--force-large', action='store_true',
                        help='Force 96x96 watermark size')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Enable verbose output')
    parser.add_argument('-q', '--quiet', action='store_true',
                        help='Suppress all output except errors')
    parser.add_argument('-b', '--banner', action='store_true',
                        help='Show ASCII banner')
    parser.add_argument('--add', action='store_true',
                        help='Add watermark instead of removing (for testing)')
    parser.add_argument('--alpha-small', 
                        help='Path to custom 48x48 alpha map')
    parser.add_argument('--alpha-large',
                        help='Path to custom 96x96 alpha map')
    parser.add_argument('--logo-value', type=float, default=255.0,
                        help='Logo pixel value (default: 255 for white)')
    parser.add_argument('--check', action='store_true',
                        help='Check alpha map files and exit')
    parser.add_argument('--version', action='version', 
                        version='Gemini Watermark Remover v1.0.0 (Python)')
    
    args = parser.parse_args()
    
    # Show banner
    if args.banner and not args.quiet:
        print_banner()
    
    # Check mode
    if args.check:
        check_alpha_maps()
        sys.exit(0)
    
    # Verbose mode (unless quiet)
    verbose = args.verbose and not args.quiet
    
    # Determine force size
    force_size = None
    if args.force_small:
        force_size = 'small'
    elif args.force_large:
        force_size = 'large'
    
    # Custom alpha map paths
    alpha_small_path = Path(args.alpha_small) if args.alpha_small else None
    alpha_large_path = Path(args.alpha_large) if args.alpha_large else None
    
    # Check default alpha maps exist (if not using custom)
    if alpha_small_path is None and not DEFAULT_ALPHA_MAP_SMALL_PATH.exists():
        print(f"Error: Alpha map not found: {DEFAULT_ALPHA_MAP_SMALL_PATH}")
        print("Please ensure bg_48.png is in the same directory as this script.")
        print("Or specify a custom path with --alpha-small")
        sys.exit(1)
    
    if alpha_large_path is None and not DEFAULT_ALPHA_MAP_LARGE_PATH.exists():
        print(f"Error: Alpha map not found: {DEFAULT_ALPHA_MAP_LARGE_PATH}")
        print("Please ensure bg_96.png is in the same directory as this script.")
        print("Or specify a custom path with --alpha-large")
        sys.exit(1)
    
    # Initialize remover
    try:
        remover = GeminiWatermarkRemover(
            alpha_map_small_path=alpha_small_path,
            alpha_map_large_path=alpha_large_path,
            logo_value=args.logo_value
        )
    except Exception as e:
        print(f"Error initializing: {e}")
        sys.exit(1)
    
    # Determine input/output
    if args.simple_input:
        # Simple mode - edit in place
        if not Path(args.simple_input).exists():
            print(f"Error: File not found: {args.simple_input}")
            sys.exit(1)
        
        try:
            process_single_file(
                args.simple_input, None, remover, 
                force_size, verbose, args.add
            )
        except Exception as e:
            print(f"Error: {e}")
            sys.exit(1)
            
    elif args.input:
        input_path = Path(args.input)
        
        if not input_path.exists():
            print(f"Error: Input not found: {args.input}")
            sys.exit(1)
        
        if input_path.is_dir():
            # Directory mode
            if not args.output:
                print("Error: Output directory required for batch processing (-o)")
                sys.exit(1)
            try:
                process_directory(
                    args.input, args.output, remover, 
                    force_size, verbose, args.add
                )
            except Exception as e:
                print(f"Error: {e}")
                sys.exit(1)
        else:
            # Single file mode
            output = args.output if args.output else args.input
            try:
                process_single_file(
                    args.input, output, remover, 
                    force_size, verbose, args.add
                )
            except Exception as e:
                print(f"Error: {e}")
                sys.exit(1)
    else:
        # No input provided - show help
        parser.print_help()
        print("\n" + "="*60)
        check_alpha_maps()
        sys.exit(0)
    
    if not args.quiet:
        print("Done!")


# ============================================================================
# ENTRY POINT
# ============================================================================

if __name__ == '__main__':
    main()