#!/usr/bin/env python3
"""Real test script for warmup functionality - uses actual client and files."""

import sys
from pathlib import Path
from sweagent.run.batch_instances import InstancesFromFile
from subagent_generator import SubagentGenerator
from tools import ToolArchive


def main():
    """Test warmup with real client and real files."""
    
    try:
        from warmup import WarmupEngine, WarmupConfig
        import openai
        
        # Use the real project paths
        project_root = Path(__file__).parent.parent
        prompt_dir = project_root / "tool_gen" / "prompts"
        template_dir = project_root / "tool_gen" / "templates_qwen3_coder_30b"
        output_dir = project_root / "tool_gen" / "test_warmup_output"
        
        print(f"Using real project paths:")
        print(f"  Project root: {project_root}")
        print(f"  Prompt dir: {prompt_dir}")
        print(f"  Template dir: {template_dir}")
        print(f"  Output dir: {output_dir}")
        
        # Verify real files exist
        agent_template = template_dir / "agent.yaml"
        subagent_template = template_dir / "subagent.yaml"
        
        # Create output directory
        output_dir.mkdir(exist_ok=True)
        
        # Create config using real paths
        config = WarmupConfig(
            warmup_iterations=5,
            output_dir=output_dir,
            prompt_dir=prompt_dir,
            template_dir=template_dir,
            model_name="openai/deepseek-v3",
            num_workers=1
        )
        
        # Create real OpenAI client with deepseek model
        client = openai.OpenAI(
            api_key="dummy-key",  # Will be overridden by api_base
            base_url="http://34.66.144.59:8001/v1"
        )
        
        print(f"Created OpenAI client with base URL: {client.base_url}")
        
        # Create engine with config
        engine = WarmupEngine(
            client=client,
            config=config
        )
        
        # Create a real tool to test with
        subagent_generator = SubagentGenerator(
            client=client,
            prompt_dir=prompt_dir,
            prompt_filename="generate_subagent_parts_prompt.txt",
            template_dir=template_dir,
            output_dir=output_dir,
            model_name="openai/deepseek-v3",
            temperature=0.5,
        )

        archive = ToolArchive(output_dir=output_dir / "subagent_tool_archive")
        test_tool = subagent_generator.generate_new_tool(archive, exp_num=0)
        test_tool.write_to_file()
        print(f"Created original tool: {test_tool.name}")
        
        # Create real instances config using the same path from your config
        instances_config = InstancesFromFile(
            path="swegym/swe_gym_lite_filtered.jsonl",
            slice=":",
            shuffle=True
        )

        # This will run the full warmup process with real instances
        result = engine.run(
            tool=test_tool,
            instances_config=instances_config,
            base_agent_template=agent_template,
            base_subagent_template=subagent_template
        )
        
    except Exception as e:
        print(f"test failed: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()
