#!/usr/bin/env python3
import json
import dotenv
from pathlib import Path
from typing import cast, Dict, Any, List
from sweagent.run.batch_instances import SWEBenchInstances
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
from sweagent.utils.serialization import merge_nested_dicts  # type: ignore
import yaml

# Load environment variables from .env file
dotenv.load_dotenv()


# Use paths relative to SWE-agent directory (where script will be run from)
output_dir = Path("tool_gen/test_output")
agent_config_path = Path("tool_gen/templates_v3/agent.yaml")

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

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

config = RunBatchConfig(
    #instances = SWEBenchInstances(
    #    subset="verified",
    #    split="test",
    #    slice="0:1"
    #),
    instances=InstancesFromFile(
        path=Path("swegym/swe_gym_lite_filtered.jsonl")
    ),
    agent=agent_config,
    output_dir=output_dir,
    num_workers=10,
    redo_existing=True,
)

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