#!/usr/bin/env python3

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import pybullet as p
from src.config_loader import ConfigLoader
from src.simulation.environment import Environment

def test_environment():
    print("Testing Environment...")
    
    # Initialize pybullet in DIRECT mode for testing
    physics_client = p.connect(p.DIRECT)
    
    try:
        # Test environment initialization
        config_loader = ConfigLoader()
        environment = Environment(config_loader, physics_client)
        print("✓ Environment initialization successful")
        
        # Test environment setup
        env_objects = environment.setup_environment()
        print("✓ Environment setup successful")
        
        # Verify all objects were created
        assert 'ground_plane_id' in env_objects, "Ground plane should be created"
        assert 'wall_parts' in env_objects, "Wall parts should be created"
        assert 'transport_object_id' in env_objects, "Transport object should be created"
        print(f"✓ Created objects: ground={env_objects['ground_plane_id']}, wall_parts={len(env_objects['wall_parts'])}, object={env_objects['transport_object_id']}")
        
        # Test transport object pose
        obj_pos, obj_orient = environment.get_transport_object_pose()
        assert len(obj_pos) == 3, "Object position should have 3 coordinates"
        assert len(obj_orient) == 4, "Object orientation should have 4 quaternion values"
        print(f"✓ Transport object pose: pos={obj_pos}, orient={obj_orient}")
        
        # Test setting object pose
        new_pos = [1.5, 0.0, 0.5]
        environment.set_transport_object_pose(new_pos)
        updated_pos, _ = environment.get_transport_object_pose()
        assert abs(updated_pos[0] - new_pos[0]) < 0.01, "Object position should be updated"
        print("✓ Transport object pose setting works")
        
        # Test hole info
        hole_info = environment.get_hole_info()
        assert 'center' in hole_info, "Hole info should include center"
        assert 'dimensions' in hole_info, "Hole info should include dimensions"
        assert 'shape' in hole_info, "Hole info should include shape"
        print(f"✓ Hole info: center={hole_info['center']}, dims={hole_info['dimensions']}")
        
        # Test wall info
        wall_info = environment.get_wall_info()
        assert 'position' in wall_info, "Wall info should include position"
        assert 'dimensions' in wall_info, "Wall info should include dimensions"
        print(f"✓ Wall info: pos={wall_info['position']}, dims={wall_info['dimensions']}")
        
        # Test environment reset
        environment.reset_environment()
        reset_pos, _ = environment.get_transport_object_pose()
        expected_pos = config_loader.get_transport_object_config()['initial_position']
        assert abs(reset_pos[0] - expected_pos[0]) < 0.01, "Object should be reset to initial position"
        print("✓ Environment reset successful")
        
        # Test collision detection
        ground_id = env_objects['ground_plane_id']
        transport_id = env_objects['transport_object_id']
        has_collision = environment.check_collision(transport_id, ground_id)
        print(f"✓ Collision detection: {has_collision}")
        
        # Test cleanup
        environment.cleanup()
        print("✓ Environment cleanup successful")
        
        print("\n✅ All Environment tests passed!")
        
        # Print environment information
        print(f"\nEnvironment Information:")
        print(f"- Wall parts created: {len(env_objects['wall_parts'])}")
        print(f"- Transport object mass: {config_loader.get_transport_object_config()['mass']} kg")
        print(f"- Hole dimensions: {hole_info['dimensions']}")
        print(f"- Wall dimensions: {wall_info['dimensions']}")
        
    except Exception as e:
        print(f"❌ Test failed: {e}")
        raise
    finally:
        # Cleanup
        p.disconnect(physics_client)
        print("✓ PyBullet disconnected")

if __name__ == "__main__":
    test_environment()