import argparse

DEFAULT_CATEGORIES = ['Animal', 'Financial', 'Privacy', 'Self-Harm', 'Violence']
AVAILABLE_CATEGORIES = ['Animal', 'Financial', 'Privacy', 'Self-Harm', 'Violence']

DEFAULT_CATEGORIES = ['bomb_explosive', 'drugs', 'firearms_weapons', 'hack_information', 'kill_someone', 'social_violence', 'suicide']
AVAILABLE_CATEGORIES = ['bomb_explosive', 'drugs', 'firearms_weapons', 'hack_information', 'kill_someone', 'social_violence', 'suicide']

def parse_categories(args):
    """Validate that all specified categories are available"""

    def validate_categories(categories):
        invalid_categories = [cat for cat in categories if cat not in AVAILABLE_CATEGORIES]
        if invalid_categories:
            print(f"Error: Invalid categories specified: {', '.join(invalid_categories)}")
            print(f"Available categories: {', '.join(AVAILABLE_CATEGORIES)}")
            return False
        return True

    if args.categories:
        # Parse comma-separated categories and strip whitespace
        category_list = [cat.strip() for cat in args.categories.split(',')]
        # Validate categories
        if not validate_categories(category_list):
            return category_list
    else:
        print("No categories specified, using default categories")
        category_list = DEFAULT_CATEGORIES
    return category_list


def CS_DJ_parser():
    
    parser = argparse.ArgumentParser(description="CS-DJ Configuration")
    parser.add_argument("--src_dir", type=str, default='../LLaVA-CC3M-Pretrain-595K/', help="Image directory")
    parser.add_argument("--save_embeding_path", type=str, default="./image_embedding", help="Path to save the embedding.")
    parser.add_argument("--save_map_path", type=str, default="./image_map", help="Path to save the jailbreak image map.")
    # parser.add_argument("--jailbreak_folder_path", type=str, default='./instructions', help="Jailbreak raw query set.")
    parser.add_argument("--jailbreak_folder_path", type=str, default='./AdvBench-M', help="Jailbreak raw query set.")
    parser.add_argument("--jailbreak_response_save_path", type=str, default='./results', help="Path to save the jailbreak response.")
    parser.add_argument("--distraction_image_save_path", type=str, default='./distraction_images', help="Path to save the jailbreak image.")
    parser.add_argument("--strategy_name", "-s", type=str, default='BSD', help="Strategy name.")
    
    # main set up
    parser.add_argument("--seed", type=int, default=2025, help="Random seed for selecting images.")
    parser.add_argument("--num_images", type=int, default=10000, help="Number of images to select.")
    parser.add_argument("--object_model", type=str, default='gpt-4o-mini', choices=['qwen2.5-vl-7b', 'qwen2.5-vl-32b', 'claude-sonnet-4','claude-haiku-3.5','gpt-4.1','gpt-4o-mini', 'gpt-4o', 'gpt-4-vision-preview', 'gpt-4.1-mini', 'llava-ov', 'internvl3-8b', 'internvl3-14b', 'internvl3-38b', 'gemini-2.5-flash', 'gemini-2.5-pro'], help="Jialbreak object model.")
    parser.add_argument("--openai_api_key", type=str, default='', help="Openai API key.")
    parser.add_argument("--anthropic_api_key", type=str, default='', help="Anthropic API key.")
    parser.add_argument("--gemini_api_key", type=str, default='', help="Gemini API key.")

    parser.add_argument("-c", "--categories", type=str, default='bomb_explosive,drugs,firearms_weapons,hack_information,kill_someone,social_violence,suicide')
    parser.add_argument("--advbench_path", type=str, default='../AdvBench-M', help="AdvBenchM data path.")
    return parser