#!/usr/bin/env python3
"""
Simple test script to validate the Container Yard environment locally.
This doesn't require OpenAI API or a running server.
"""

import sys
sys.path.insert(0, '.')

from server.Container_Yard_environment import ContainerYardEnvironment
from models import ContainerYardAction

def test_task(task_name: str, num_steps: int = 100):
    """Test a single task."""
    print(f"\n{'='*60}")
    print(f"Testing {task_name.upper()} task")
    print(f"{'='*60}")
    
    env = ContainerYardEnvironment(task_name=task_name)
    obs = env.reset()
    
    print(f"Initial state:")
    print(f"  - Containers: {obs.total_containers}")
    print(f"  - Stacks: {obs.num_stacks}")
    print(f"  - Max height: {obs.max_stack_height}")
    
    step_count = 0
    total_reward = 0.0
    
    while not obs.done and step_count < num_steps:
        step_count += 1
        
        # Simple greedy strategy: place in stack with fewest containers
        stack_sizes = [len(stack) for stack in obs.stacks]
        best_stack = min(range(len(stack_sizes)), key=lambda i: stack_sizes[i])
        
        action = ContainerYardAction(stack_index=best_stack)
        obs = env.step(action)
        total_reward += obs.reward
        
        if step_count <= 3 or obs.done:
            print(f"Step {step_count}: Placed C{obs.current_container_id-1 if obs.current_container_id > 0 else 'X'} "
                  f"in stack {best_stack}, reward={obs.reward:.2f}, "
                  f"rehandles={obs.rehandles_so_far}")
    
    efficiency = 1.0 - (obs.rehandles_so_far / obs.total_containers)
    
    print(f"\nFinal results:")
    print(f"  - Steps taken: {step_count}")
    print(f"  - Total reward: {total_reward:.2f}")
    print(f"  - Total rehandles: {obs.rehandles_so_far}")
    print(f"  - Efficiency score: {efficiency:.2f}")
    print(f"  - Success: {obs.rehandles_so_far <= 3 * (len(obs.stacks) // 5)}")
    
    return {
        'task': task_name,
        'steps': step_count,
        'rehandles': obs.rehandles_so_far,
        'efficiency': efficiency,
        'total_reward': total_reward,
    }

if __name__ == "__main__":
    print("Container Yard Environment - Local Test")
    
    results = []
    for task in ['easy', 'medium', 'hard']:
        try:
            result = test_task(task)
            results.append(result)
        except Exception as e:
            print(f"ERROR in {task} task: {e}")
            import traceback
            traceback.print_exc()
    
    print(f"\n{'='*60}")
    print("Summary")
    print(f"{'='*60}")
    for r in results:
        print(f"{r['task']:8} | Steps: {r['steps']:2} | Rehandles: {r['rehandles']:2} | Efficiency: {r['efficiency']:.2f}")