"""
Class to set up a 6-wheeled lunar cruiser, loaded from an openplx file
"""

import os

from agxPythonModules.utils.environment import application, root, simulation
from agxPythonModules.utils.callbacks import KeyboardCallback, GamepadCallback

from openplx.Physics.Signals import RealInputSignal
import agx
import agxSDK
from agxOpenPLX import OptParams

class VehicleManeuvering:
    """
    Class to manage vehicle maneuvering using input signals.
    Sets up keyboard and gamepad inputs:
        Keyboard:
            - Turn left/right:  left/right arrows
            - Drive forward:    up arrow
            - Reverse:          down arrow

        Gamepad (Xbox like):
            - Turn left/right:  left stick horizontal
            - Drive forward:    right trigger
            - Reverse:          left trigger
    """

    def __init__(self, lunar_cruiser):
        """
        Initialize the vehicle maneuvering class with input signals control.

        :param lunar_cruiser: LunarCruiser class, representing vehicle to be controlled
        """
        self.target_velocity = 0  # Default target velocity
        self.amplitude = 2  # Speed amplitude for the motors
        self.steering_angle = 0
        self.lunar_cruiser = lunar_cruiser

    def forward(self):
        """
        Set target velocity to start forward movement.
        """
        self.target_velocity = self.amplitude
        self.lunar_cruiser.set_motor_velocity(self.target_velocity)

    def backward(self):
        """
        Set target velocity to start backward movement.
        """
        self.target_velocity = -self.amplitude
        self.lunar_cruiser.set_motor_velocity(self.target_velocity)

    def left(self):
        """ 
        Handle steering left.
        """
        if self.steering_angle < 0.6:
            self.steering_angle += 0.01
        self.lunar_cruiser.set_steer_angle(self.steering_angle)

    def right(self):
        """
        Handle steering right.
        """
        if self.steering_angle > -0.6:
            self.steering_angle -= 0.01
        self.lunar_cruiser.set_steer_angle(self.steering_angle)

    def stop(self):
        """
        Stop by setting target velocity to 0.
        """
        self.target_velocity = 0
        self.lunar_cruiser.set_motor_velocity(self.target_velocity)

    def setup_keyboard_listener(self):
        """
        Set up keyboard listeners for forward, backward and steering controls.
        """
        KeyboardCallback.bind(
            name="forward",
            key=agxSDK.GuiEventListener.KEY_Up,
            callback=lambda key: (self.forward() if key.down else self.stop()),
        )
        KeyboardCallback.bind(
            name="backward",
            key=agxSDK.GuiEventListener.KEY_Down,
            callback=lambda key: (self.backward() if key.down else self.stop()),
        )
        KeyboardCallback.bind(
            name="left",
            key=agxSDK.GuiEventListener.KEY_Left,
            callback=lambda key: self.left(),
        )
        KeyboardCallback.bind(
            name="right",
            key=agxSDK.GuiEventListener.KEY_Right,
            callback=lambda key: self.right(),
        )
    
    def setup_gamepad_listener(self):
        """
        Set up gamepad listeners for forward, backward and steering controls.
        """
        
        def steer_callback(value):
            if value < -0.2:
                self.left()
            elif value > 0.2:
                self.right()
        
        def drive_callback(data):
            if data.value > 0.2:
                if data.name == "reverse":
                    self.backward()
                else:
                    self.forward()
                    
            elif data.delta != 0:
                self.stop()

        GamepadCallback.bind(
            name="LeftRight",
            axis=GamepadCallback.Axis.LeftHorizontal,
            callback=lambda data: steer_callback(data.value),
        )

        GamepadCallback.bind(
            name="forward",
            axis=GamepadCallback.Axis.RightTrigger,
            callback=drive_callback,
        )

        GamepadCallback.bind(
            name="reverse",
            axis=GamepadCallback.Axis.LeftTrigger,
            callback=drive_callback,
        )


class AlgoryxLunarCruiser:
    """
    Class to setup lunar cruiser with 6 wheels from openplx file
    """
    def __init__(self, use_drivetrain=False):
        """
        Initialize lunar cruiser. By default the vehicle uses hinge motors on each wheel to 
        drive it forward.

        :param use_drivetrain: Set to true to use drive train with differentials and a single 
                               velocity motor. Default false.
        """
        file_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        if use_drivetrain:
            model_path = os.path.join(file_dir, "CruiserOneMotor.openplx")
        else:
            model_path = os.path.join(file_dir, "CruiserSixMotors.openplx")
        
        result = application().loadOpenPlxFile(model_path, root())
        openplx_scene = result.scene()
        self.assembly = result.assembly()

        self.signal_interface = openplx_scene.getObject("signal_interface")
        self.input_queue = result.getInputSignalQueue()
        self.inputs = self.signal_interface.getInputs()
        self.use_drivetrain = use_drivetrain
        self.model_name = result.scene().getType().getName()

        # Add VehicleManeuvering to be able to add keyboard and gamepad controls
        self.vehicle_maneuvering = VehicleManeuvering(self)
   
    def set_steer_angle(self, angle):
        """
        Set steering angle of front wheel steering by sending openplx-signal.

        :param angle: Steering angle
        """
        self.input_queue.send(RealInputSignal.create(angle, self.inputs["front_steering_angle_input"]))
 
    def set_motor_velocity(self, target_angular_velocity):
        """
        Set target target angular velocity on motor/motors.

        :param target_angular_velocity: Target angular velocity for motor/motors.
        """
        if self.use_drivetrain:
            self.input_queue.send(RealInputSignal.create(target_angular_velocity, 
                                                         self.inputs["motor_vel_input"]))
        else:
            self.input_queue.send(RealInputSignal.create(target_angular_velocity,
                                                         self.inputs["front_left_vel_input"]))
            self.input_queue.send(RealInputSignal.create(target_angular_velocity,
                                                         self.inputs["front_right_vel_input"]))
            self.input_queue.send(RealInputSignal.create(target_angular_velocity,
                                                         self.inputs["rear_left_vel_input"]))
            self.input_queue.send(RealInputSignal.create(target_angular_velocity,
                                                         self.inputs["rear_right_vel_input"]))
            self.input_queue.send(RealInputSignal.create(target_angular_velocity,
                                                         self.inputs["mid_left_vel_input"]))
            self.input_queue.send(RealInputSignal.create(target_angular_velocity,
                                                         self.inputs["mid_right_vel_input"]))
    
    def setup_wheel_ground_contact_material(self, ground_material):
        """
        Untility method to helt set up contact material between vehicle wheels and a ground material.
        
        :param ground_material: Material for the ground the wheels drive on
        """
        # Create contact material
        wheel_material = simulation().getMaterialManager().getMaterial("WheelMaterial")
        assert wheel_material, "Could not find wheel material."
        tire_ground_contact_material = (
            simulation().getMaterialManager().getOrCreateContactMaterial(wheel_material, ground_material)
        )
        simulation().add(tire_ground_contact_material)

        # Set contact material's friction coeffs
        tire_ground_contact_material.setFrictionCoefficient(1.0, agx.ContactMaterial.PRIMARY_DIRECTION)
        tire_ground_contact_material.setFrictionCoefficient(1.0, agx.ContactMaterial.SECONDARY_DIRECTION)

        # Set contact material restitution
        tire_ground_contact_material.setRestitution(0.0)

        # Set fricton model
        tire_ground_friction = agx.ScaleBoxFrictionModel()
        tire_ground_friction.setSolveType(agx.FrictionModel.DIRECT_AND_ITERATIVE)
        tire_ground_contact_material.setFrictionModel(tire_ground_friction)
    
    def enable_keyboard_gamepad_control(self):
        """
        Enables gamepad and keyboard controls.
        """
        self.vehicle_maneuvering.setup_keyboard_listener()
        self.vehicle_maneuvering.setup_gamepad_listener()
