import os
import argparse
from tqdm import tqdm
from armor import ArmorGenerator, ArmorConfig

def batch_armor_images(input_dir, output_dir, strength=2.5):
    """
    Applies Ghostprint armor to a directory of images.

    Args:
        input_dir (str): Directory containing the original images.
        output_dir (str): Directory to save the armored images.
        strength (float): The strength of the armor to apply (1.0 to 5.0).
    """
    if not os.path.isdir(input_dir):
        print(f"Error: Input directory '{input_dir}' not found.")
        return

    os.makedirs(output_dir, exist_ok=True)

    # Initialize the ArmorGenerator
    # Using default ArmorConfig but can be customized if needed
    armor_generator = ArmorGenerator()
    armor_config = ArmorConfig()

    image_files = sorted([f for f in os.listdir(input_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
    
    if not image_files:
        print(f"No images found in '{input_dir}'.")
        return

    print(f"Applying armor with strength {strength} to {len(image_files)} images...")

    for filename in tqdm(image_files, desc="Armoring images"):
        input_path = os.path.join(input_dir, filename)
        output_path = os.path.join(output_dir, filename)

        try:
            armor_generator.apply_to_image(
                image_path=input_path,
                out_path=output_path,
                cfg=armor_config,
                strength=strength,
                return_metrics=False # We don't need metrics for batch processing
            )
        except Exception as e:
            print(f"\nFailed to process {filename}: {e}")

    print("Batch armoring complete.")
    print(f"Armored images saved in '{output_dir}'.")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Batch apply Ghostprint armor to images.")
    parser.add_argument(
        '--input-dir', 
        type=str, 
        default='placeholder_images',
        help='Directory containing the original images.'
    )
    parser.add_argument(
        '--output-dir', 
        type=str, 
        default='armored_images',
        help='Directory to save the armored images.'
    )
    parser.add_argument(
        '--strength', 
        type=float, 
        default=2.5,
        help='Armor strength (e.g., 1.0 to 5.0). 2.5 is a good starting point.'
    )
    args = parser.parse_args()

    batch_armor_images(args.input_dir, args.output_dir, args.strength)