using AGXUnity;
using agxUtil;
using UnityEngine;
using AGXUnity.Utils;

#if UNITY_EDITOR
using UnityEditor;
using AGXUnityEditor;
#endif

public class MachineResetter : ScriptComponent
{
  private agx.RigidBody m_mainBody;
  private agx.AffineMatrix4x4 m_mainTransform;

  private agxSDK.Collection m_collection;
  private agxUtil.ConstraintPositionVector m_constraintPositions;

  private agxSDK.Assembly m_jumpAssembly;

  public int NumBodies => m_collection?.getRigidBodies().Count ?? 0;
  public int NumConstraints => m_collection?.getConstraints().Count ?? 0;
  public Vector3 MainBodyPos => m_mainTransform.getTranslate().ToHandedVector3();
  public Vector3 MainBodyRot => m_mainTransform.getRotate().ToHandedQuaternion().eulerAngles;

  protected override bool Initialize()
  {
    // Find all affected components
    foreach ( var sc in GetComponentsInChildren<ScriptComponent>() )
      if ( sc != this )
        sc.GetInitialized();

    SaveCurrentPose();

    return base.Initialize();
  }

  [InvokableInInspector]
  public void SaveCurrentPose()
  {
    m_collection = new agxSDK.Collection();
    m_constraintPositions = new agxUtil.ConstraintPositionVector();
    m_jumpAssembly = new agxSDK.Assembly();

    var bodies = GetComponentsInChildren<RigidBody>();
    var constraints = GetComponentsInChildren<Constraint>();
    var tires = GetComponentsInChildren<AGXUnity.Model.TwoBodyTire>();

    // Add initial native components properties to native collections
    foreach ( var bod in bodies ) {
      m_collection.add( bod.GetInitialized().Native );
      m_jumpAssembly.add( bod.Native );
    }
    foreach ( var constraint in constraints ) {
      if ( constraint.isActiveAndEnabled ) {
        var init = constraint.GetInitialized();
        m_collection.add( init.Native );
        m_constraintPositions.Add( new agxUtil.ConstraintPosition( constraint.Native, constraint.GetCurrentAngle() ) );
      }
    }
    foreach ( var tire in tires )
      m_collection.add( tire.GetInitialized().Native );

    // save the initial transform of the main body
    m_mainBody = bodies[ 0 ].Native;
    m_mainTransform = m_mainBody.getTransform();
  }

  public void LoadSavedPose()
  {
    // Apply transform for main body
    agxUtilSWIG.jumpRequest( m_jumpAssembly, m_mainBody, m_mainTransform );

    // Compute and apply new configuration
    var reconfRequest = new agxUtil.ReconfigureRequest();
    agxUtil.BodyTransformVector result = new agxUtil.BodyTransformVector();
    reconfRequest.computeTransforms( m_collection, m_mainBody, m_constraintPositions, result );
    reconfRequest.applyTransforms( m_collection, m_constraintPositions, result );

    // Applying new transforms to a machine might cause enabled LockControllers to be violated.
    // Here, we set the locked position for all enabled locks after applying the transformation 
    // avoid any potential violations.
    foreach ( var constraint in GetComponentsInChildren<Constraint>() ) {
      var transLock = constraint.GetController<LockController>(Constraint.ControllerType.Translational);
      if ( transLock != null )
        if ( transLock.Enable )
          transLock.Position = constraint.GetCurrentAngle( Constraint.ControllerType.Rotational );

      var rotLock = constraint.GetController<LockController>(Constraint.ControllerType.Rotational);
      if ( rotLock != null )
        if ( rotLock.Enable )
          rotLock.Position = constraint.GetCurrentAngle( Constraint.ControllerType.Rotational );
    }
  }

  // Update is called once per frame
  void Update()
  {
    if ( Input.GetKeyDown( KeyCode.T ) )
      SaveCurrentPose();
    if ( Input.GetKeyDown( KeyCode.R ) )
      LoadSavedPose();
  }
}

#if UNITY_EDITOR
[CustomEditor( typeof( MachineResetter ) )]
class MachineResetterEditor : Editor
{
  public override void OnInspectorGUI()
  {
    var resetter = (MachineResetter)target;
    using var dis = new EditorGUI.DisabledScope(!Application.isPlaying);
    GUILayout.Label( "Currently Saved Pose:" );
    using ( new InspectorGUI.IndentScope() ) {
      EditorGUILayout.LabelField( "# saved bodies:", $"{resetter.NumBodies}" );
      EditorGUILayout.LabelField( "# saved constraints:", $"{resetter.NumConstraints}" );
      var pos = resetter.MainBodyPos;
      EditorGUILayout.LabelField( "Saved position:", $"({pos.x:f2}, {pos.y:f2}, {pos.z:f2})" );
      var rot = resetter.MainBodyRot;
      EditorGUILayout.LabelField( "Saved rotation:", $"({rot.x:f1}, {rot.y:f1}, {rot.z:f1})" );
    }
    if ( GUILayout.Button( "Load Saved Pose" ) ) {
      resetter.LoadSavedPose();
    }
    if ( GUILayout.Button( "Save Current Pose" ) )
      resetter.SaveCurrentPose();
  }
}
#endif
