#!/usr/bin/env python3
"""
Example usage of the updated AgentariumTrainer with AgentOrchestrator.

This demonstrates the new Ray-based task dispatch system instead of HTTP API.
"""

import os
import sys
import time
from pathlib import Path

# Add the project root to the path
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from agentlightning.agentarium.src.orchestrator import AgentOrchestrator
from examples.custom_calc_x.calc_agent_ray import calcagent_rollout

def main():
    """Demonstrate the new Ray-based AgentOrchestrator system."""
    
    print("=== AgentariumTrainer with Ray-based AgentOrchestrator ===\n")
    
    # Initialize the orchestrator
    orchestrator = AgentOrchestrator()
    
    print("1. Initializing Ray cluster and registering calc agent...")
    orchestrator.init_ray()
    orchestrator.register_agent_type("calc_agent", calcagent_rollout, "https://api.openai.com/v1")
    print(f"   Registered agent types: {list(orchestrator.agent_types.keys())}\n")
    
    # Define calculation tasks
    tasks = [
        {"question": "What is 15 + 27 * 3?", "result": "96"},
        {"question": "Calculate the area of a circle with radius 8 meters", "result": "201.06"},
        {"question": "What's the square root of 144?", "result": "12"},
        {"question": "Find the factorial of 5", "result": "120"},
        {"question": "Convert 100 degrees Fahrenheit to Celsius", "result": "37.78"},
    ]
    
    # Submit tasks using the agent-lightning framework interface
    print("2. Setting up data and server...")
    data = {
        "question": [task["question"] for task in tasks],
        "result": [task["result"] for task in tasks]
    }
    
    # Simulate the trainer interface
    orchestrator.initialize_batch_processing(data, server_addresses=[], is_train=False)
    
    print("3. Running until all tasks finished...")
    orchestrator.wait_for_batch_completion(verbose=True)
    
    print("4. Getting test metrics...")
    test_metrics = orchestrator.collect_validation_metrics()
    print(f"   Test metrics: {test_metrics}")
    
    print("5. Clearing data and server...")
    orchestrator.reset_orchestrator_state()
    
    # Demonstrate training data batch generation
    print("\n6. Demonstrating training data batch generation...")
    
    # Set up for training
    orchestrator.initialize_batch_processing(data, server_addresses=[], is_train=True)
    orchestrator.wait_for_batch_completion(verbose=False)
    
    # Mock device for demonstration
    import torch
    device = torch.device("cpu")
    
    try:
        data_proto, metrics = orchestrator.generate_training_batch(
            max_prompt_length=128,
            max_response_length=64,
            device=device
        )
        print(f"   Generated batch with {data_proto.batch.batch_size} samples")
        print(f"   Batch metrics: {metrics}")
    except Exception as e:
        print(f"   Error generating batch: {e}")
    
    orchestrator.reset_orchestrator_state()
    
    print("\n7. Shutting down orchestrator...")
    orchestrator.shutdown()
    print("   Done!")


def demonstrate_direct_ray_usage():
    """
    Demonstrate using the Ray functions directly without the trainer interface.
    """
    print("\n=== Direct Ray Usage Demo ===\n")
    
    import ray
    from examples.custom_calc_x.calc_agent_ray import calcagent_rollout
    
    # Initialize Ray
    if not ray.is_initialized():
        ray.init()
    
    # Submit tasks directly to Ray
    tasks = [
        {"question": "What is 2^10?", "result": "1024"},
        {"question": "Calculate 25% of 400", "result": "100"},
        {"question": "What's the area of a triangle with base 10 and height 6?", "result": "30"},
    ]
    
    print("Submitting tasks directly to Ray...")
    futures = []
    for task in tasks:
        future = calcagent_rollout.remote(task, "https://api.openai.com/v1")
        futures.append(future)
        print(f"   Submitted: {task['question']}")
    
    print("\nWaiting for results...")
    results = ray.get(futures)
    
    print("\nResults:")
    for i, result in enumerate(results):
        print(f"   Task {i+1}: {result.question}")
        print(f"   Answer: {result.answer}")
        print(f"   Reward: {result.reward}")
        print(f"   Ground Truth: {result.ground_truth}")
        print()
    
    ray.shutdown()
    print("Ray shutdown complete!")


if __name__ == "__main__":
    # Run the main demo
    main()
    
    # Uncomment to run the direct Ray usage demo
    # demonstrate_direct_ray_usage() 