#!/usr/bin/env python3
"""
Convert AdvBench-M JSON format to instructions format.

AdvBench-M format:
{
    "scenario": "category_name",
    "instructions": ["prompt1", "prompt2", ...],
    ...
}

Target format:
[
    {"id": 1, "instruction": "prompt1"},
    {"id": 2, "instruction": "prompt2"},
    ...
]
"""

import json
import os
import sys
from pathlib import Path
from config import CS_DJ_parser


def convert_advbench_to_instructions(input_file, output_file=None):
    """
    Convert a single AdvBench-M JSON file to instructions format.
    
    Args:
        input_file (str): Path to input AdvBench-M JSON file
        output_file (str, optional): Path to output file. If None, auto-generate based on scenario name
    
    Returns:
        str: Path to the output file
    """
    # Read input file
    with open(input_file, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    # Extract scenario and instructions
    scenario = data.get('scenario', 'unknown')
    instructions = data.get('instructions', [])
    
    # Convert to target format
    converted_data = []
    for i, instruction in enumerate(instructions, 1):
        converted_data.append({
            "id": i,
            "instruction": instruction
        })
    
    # Determine output file path
    if output_file is None:
        # Capitalize scenario name for output filename
        # scenario_name = scenario.replace('_', ' ').title().replace(' ', '')
        output_file = f"{scenario}.json"
    
    # Write output file
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(converted_data, f, indent=4, ensure_ascii=False)
    
    print(f"Converted {len(instructions)} instructions from '{scenario}' to {output_file}")
    return output_file



def convert_from_config(args):
    """Convert AdvBench-M files based on configuration."""
    advbench_path = Path(args.advbench_path)
    jailbreak_folder_path = Path(args.jailbreak_folder_path)
    
    # Create output directory if it doesn't exist
    jailbreak_folder_path.mkdir(parents=True, exist_ok=True)
    
    # Look for JSON files in AdvBench-M directory
    instructions_dir = advbench_path / "prompts" / "eval_all_instructions"
    
    if not instructions_dir.exists():
        print(f"Error: AdvBench-M instructions directory not found at {instructions_dir}")
        return
    
    json_files = list(instructions_dir.glob("*.json"))
    
    if not json_files:
        print(f"No JSON files found in {instructions_dir}")
        return
    
    print(f"Found {len(json_files)} AdvBench-M files to convert...")
    
    for json_file in json_files:
        try:
            # Read the AdvBench-M file to get scenario name
            with open(json_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            scenario = data.get('scenario', json_file.stem)
            
            # Create output filename based on scenario
            # scenario_name = scenario.replace('_', ' ').title().replace(' ', '')
            output_file = jailbreak_folder_path / f"{scenario}.json"
            
            convert_advbench_to_instructions(str(json_file), str(output_file))
            
        except Exception as e:
            print(f"Error converting {json_file}: {e}")


def convert_single_category(args):
    """Convert a specific category from AdvBench-M."""
    advbench_path = Path(args.advbench_path)
    jailbreak_folder_path = Path(args.jailbreak_folder_path)
    category = args.category.lower()
    
    # Create output directory if it doesn't exist
    jailbreak_folder_path.mkdir(parents=True, exist_ok=True)
    
    # Look for the specific category file
    instructions_dir = advbench_path / "prompts" / "all_instructions"
    category_file = instructions_dir / f"{category}.json"
    
    if not category_file.exists():
        print(f"Error: Category file not found at {category_file}")
        # Try to find files with similar names
        possible_files = list(instructions_dir.glob(f"*{category}*.json"))
        if possible_files:
            print("Possible matches found:")
            for file in possible_files:
                print(f"  - {file.name}")
        return
    
    try:
        # Create output filename
        category_name = args.category.title()
        output_file = jailbreak_folder_path / f"{category_name}.json"
        
        convert_advbench_to_instructions(str(category_file), str(output_file))
        
    except Exception as e:
        print(f"Error converting {category_file}: {e}")


def main():
    """Main function using CS_DJ_parser configuration."""
    parser = CS_DJ_parser()
    
    # Add specific arguments for conversion script
    parser.add_argument("--convert_all", action="store_true", 
                       help="Convert all AdvBench-M files to instructions format")
    parser.add_argument("--convert_category", action="store_true",
                       help="Convert only the specified category")
    parser.add_argument("--input_file", type=str,
                       help="Specific input file to convert")
    parser.add_argument("--output_file", type=str,
                       help="Specific output file path")
    
    args = parser.parse_args()
    
    if args.input_file:
        # Single file mode
        if not os.path.exists(args.input_file):
            print(f"Error: Input file '{args.input_file}' does not exist")
            return
        
        convert_advbench_to_instructions(args.input_file, args.output_file)
        
    elif args.convert_all:
        # Convert all AdvBench-M files
        convert_from_config(args)
        
    elif args.convert_category:
        # Convert specific category
        convert_single_category(args)
        
    else:
        print("Usage:")
        print("  Convert all files:       python cvt_advbench2hades.py --convert_all")
        print("  Convert specific category: python cvt_advbench2hades.py --convert_category --category Animal")
        print("  Convert single file:     python cvt_advbench2hades.py --input_file path/to/file.json [--output_file output.json]")
        print("\nConfiguration options:")
        print(f"  AdvBench-M path: {args.advbench_path}")
        print(f"  Output directory: {args.jailbreak_folder_path}")
        print(f"  Default category: {args.categories}")


if __name__ == "__main__":
    main()
