#!/usr/bin/env python3
import json
import argparse
from pathlib import Path
from typing import cast, Dict, Any, List

from datasets import Dataset, load_dataset
from sweagent.run.run_batch import run_from_config, RunBatchConfig  # type: ignore
from sweagent.run.batch_instances import InstancesFromFile  # changed from InstancesFromHuggingFace
from sweagent.agent.agents import DefaultAgentConfig  # type: ignore
import yaml

# Parse command line arguments
parser = argparse.ArgumentParser(description='Run batch test with configurable output directory suffix')
parser.add_argument('--suffix', type=str, default='swe_gym_lite_3', 
                   help='Suffix for output directory name (default: swe_gym_lite_3)')
args = parser.parse_args()

# Use paths relative to SWE-agent directory (where script will be run from)
# use today's month and day as the output directory

import datetime
today = datetime.datetime.now().strftime("%B%d")
output_dir = Path(f"tool_gen/v3_output_{args.suffix}")

# Load agent config from generated experiment
agent_config_path = Path("/root/SWE-agent/scripts/agent_without_str.yaml")

with open(agent_config_path) as f:
    config_data = yaml.safe_load(f)


agent_config = DefaultAgentConfig.model_validate(config_data["agent"])

# Run batch with filtered SWE-smith instances from file
config = RunBatchConfig(
    instances=InstancesFromFile(
        path=Path("swe_gym_lite_filtered.jsonl"),
        filter=".*",
        slice="0:1000"
    ),
    agent=agent_config,
    output_dir=output_dir,
    num_workers=32,
    redo_existing=False,
)

print("Running batch test...")
print(f"Agent config: {agent_config_path}")
print(f"Dataset: filtered_swe_smith.jsonl")
print(f"Output directory: {output_dir}")
run_from_config(config) 