#!/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.robot.robot_loader import RobotLoader

def test_robot_loader():
    print("Testing RobotLoader...")
    
    # Initialize pybullet in DIRECT mode for testing
    physics_client = p.connect(p.DIRECT)
    
    try:
        # Test config loading
        config_loader = ConfigLoader()
        robot_loader = RobotLoader(config_loader, physics_client)
        print("✓ RobotLoader initialization successful")
        
        # Test robot loading
        robot_id = robot_loader.load_robot()
        assert robot_id is not None, "Robot ID should not be None"
        print(f"✓ Robot loaded successfully with ID: {robot_id}")
        
        # Test DOF
        dof = robot_loader.get_dof()
        print(f"✓ Robot has {dof} degrees of freedom")
        
        # Test joint names
        joint_names = robot_loader.get_joint_names()
        assert len(joint_names) == dof, "Joint names count should match DOF"
        print(f"✓ Joint names: {joint_names[:3]}..." if len(joint_names) > 3 else f"✓ Joint names: {joint_names}")
        
        # Test joint limits
        joint_limits = robot_loader.get_joint_limits()
        assert len(joint_limits) == dof, "Joint limits count should match DOF"
        print("✓ Joint limits retrieved successfully")
        
        # Test joint states
        joint_states = robot_loader.get_joint_states()
        assert len(joint_states) == dof, "Joint states count should match DOF"
        print("✓ Joint states retrieved successfully")
        
        # Test end effector pose
        ee_pos, ee_orient = robot_loader.get_end_effector_pose()
        assert len(ee_pos) == 3, "End effector position should have 3 coordinates"
        assert len(ee_orient) == 4, "End effector orientation should have 4 quaternion values"
        print(f"✓ End effector pose: pos={ee_pos}, orient={ee_orient}")
        
        # Test joint position validation
        zero_positions = [0.0] * dof
        is_valid = robot_loader.validate_joint_positions(zero_positions)
        print(f"✓ Zero positions validation: {is_valid}")
        
        # Test setting joint positions
        robot_loader.set_joint_positions(zero_positions)
        print("✓ Joint positions set successfully")
        
        # Test robot reset
        robot_loader.reset_robot()
        print("✓ Robot reset successful")
        
        # Test self-collision check
        has_collision = robot_loader.check_self_collision()
        print(f"✓ Self-collision check: {has_collision}")
        
        # Test control mode switching
        robot_loader.enable_position_control()
        print("✓ Position control mode enabled")
        
        robot_loader.enable_torque_control()
        print("✓ Torque control mode enabled")
        
        print("\n✅ All RobotLoader tests passed!")
        
        # Print robot information
        print(f"\nRobot Information:")
        print(f"- DOF: {dof}")
        print(f"- Robot ID: {robot_id}")
        print(f"- Joint count: {robot_loader.num_joints}")
        print(f"- Controllable joints: {len(robot_loader.joint_indices)}")
        print(f"- End effector index: {robot_loader.end_effector_index}")
        
    except Exception as e:
        print(f"❌ Test failed: {e}")
        raise
    finally:
        # Cleanup
        p.disconnect(physics_client)
        print("✓ PyBullet disconnected")

if __name__ == "__main__":
    test_robot_loader()