"""
Keyboard controller for Scenes/DozerOnPlane.openplx.
"""

import os
from enum import Enum

import agx
import agxOpenPLX
import agxOSG
import agxSDK
import openplx
from agxPythonModules.tools.simulation_content import SimulationContent
from agxPythonModules.utils.callbacks import KeyboardCallback as Input
from agxPythonModules.utils.environment import (
    application,
    init_app,
    root,
    simulation,
)
from openplx.Physics.Signals import IntInputSignal, RealInputSignal


class Position(Enum):
    LEFT = 0
    RIGHT = 1


class DriveControlMode(Enum):
    MOTOR_VELOCITY = 0
    ENGINE_THROTTLE = 1


def model_bundle_path() -> str:
    return os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def get_required_object(scene, *names):
    for name in names:
        obj = scene.getObject(name)
        if obj:
            return obj
    raise AssertionError(f"None of the expected OpenPLX objects were found: {names}")


def get_linear_velocity_input(scene, *base_names):
    candidate_names = []
    for base_name in base_names:
        candidate_names.append(f"{base_name}.source.target_linear_velocity_input")
        candidate_names.append(f"{base_name}.target_linear_velocity_input")
        candidate_names.append(f"{base_name}.motor.target_linear_velocity_input")
    return get_required_object(scene, *candidate_names)


def get_optional_linear_velocity_input(scene, *base_names):
    candidate_names = []
    for base_name in base_names:
        candidate_names.append(f"{base_name}.source.target_linear_velocity_input")
        candidate_names.append(f"{base_name}.target_linear_velocity_input")
        candidate_names.append(f"{base_name}.motor.target_linear_velocity_input")
    for name in candidate_names:
        obj = scene.getObject(name)
        if obj:
            return obj
    return None


class BulldozerManeuver:  # pylint: disable=too-many-instance-attributes
    def __init__(
        self,
        input_queue,
        drive_input: openplx.Physics.Signals.Input,
        drive_control_mode: DriveControlMode,
        gear_input: openplx.Physics.Signals.IntInput | None,
        brake_control: dict,
        blade_control: dict,
        ripper_control: dict,
    ):
        self.input_queue = input_queue
        self.drive_input = drive_input
        self.drive_control_mode = drive_control_mode
        self.gear_input = gear_input
        self.clutch_left_input = brake_control["clutch_left_input"]
        self.clutch_right_input = brake_control["clutch_right_input"]
        self.brake_left_input = brake_control["brake_left_input"]
        self.brake_right_input = brake_control["brake_right_input"]
        self.blade_left_lift_input = blade_control["blade_left_lift_input"]
        self.blade_right_lift_input = blade_control["blade_right_lift_input"]
        self.blade_left_tilt_input = blade_control["blade_left_tilt_input"]
        self.blade_right_tilt_input = blade_control["blade_right_tilt_input"]
        self.ripper_left_lift_input = ripper_control["ripper_left_lift_input"]
        self.ripper_right_lift_input = ripper_control["ripper_right_lift_input"]

        self.drive_target = 0.0
        self.motor_velocity_amplitude = 418.0
        self.engine_throttle_amplitude = 0.85
        self.target_gear = 0
        self.clutch_target_engagement_fraction = {
            Position.LEFT: 1.0,
            Position.RIGHT: 1.0,
        }
        self.brake_target_engagement_fraction = {
            Position.LEFT: 0.0,
            Position.RIGHT: 0.0,
        }
        self.blade_lift_target_velocity = 0.0
        self.blade_lift_velocity_amplitude = 0.2
        self.blade_tilt_target_velocity = 0.0
        self.blade_tilt_velocity_amplitude = 0.2
        self.ripper_lift_target_velocity = 0.0
        self.ripper_lift_velocity_amplitude = 0.1

    def send_signals(self):
        self.input_queue.send(
            RealInputSignal.create(self.drive_target, self.drive_input)
        )
        if self.gear_input is not None:
            self.input_queue.send(IntInputSignal.create(self.target_gear, self.gear_input))
        self.input_queue.send(
            RealInputSignal.create(
                self.clutch_target_engagement_fraction[Position.LEFT],
                self.clutch_left_input,
            )
        )
        self.input_queue.send(
            RealInputSignal.create(
                self.clutch_target_engagement_fraction[Position.RIGHT],
                self.clutch_right_input,
            )
        )
        self.input_queue.send(
            RealInputSignal.create(
                self.brake_target_engagement_fraction[Position.LEFT],
                self.brake_left_input,
            )
        )
        self.input_queue.send(
            RealInputSignal.create(
                self.brake_target_engagement_fraction[Position.RIGHT],
                self.brake_right_input,
            )
        )
        self.input_queue.send(
            RealInputSignal.create(
                self.blade_lift_target_velocity,
                self.blade_left_lift_input,
            )
        )
        self.input_queue.send(
            RealInputSignal.create(
                self.blade_lift_target_velocity,
                self.blade_right_lift_input,
            )
        )
        self.input_queue.send(
            RealInputSignal.create(
                self.blade_tilt_target_velocity,
                self.blade_left_tilt_input,
            )
        )
        self.input_queue.send(
            RealInputSignal.create(
                self.blade_tilt_target_velocity,
                self.blade_right_tilt_input,
            )
        )
        if self.ripper_left_lift_input is not None:
            self.input_queue.send(
                RealInputSignal.create(
                    self.ripper_lift_target_velocity,
                    self.ripper_left_lift_input,
                )
            )
        if self.ripper_right_lift_input is not None:
            self.input_queue.send(
                RealInputSignal.create(
                    self.ripper_lift_target_velocity,
                    self.ripper_right_lift_input,
                )
            )

    def forward(self):
        if self.drive_control_mode == DriveControlMode.MOTOR_VELOCITY:
            self.drive_target = self.motor_velocity_amplitude
        else:
            self.drive_target = self.engine_throttle_amplitude
            if self.gear_input is not None:
                self.target_gear = 1
        self.send_signals()

    def backward(self):
        if self.drive_control_mode == DriveControlMode.MOTOR_VELOCITY:
            self.drive_target = -self.motor_velocity_amplitude
        else:
            self.drive_target = self.engine_throttle_amplitude
            if self.gear_input is not None:
                self.target_gear = -1
        self.send_signals()

    def brake_left_engage(self):
        self.clutch_target_engagement_fraction[Position.LEFT] = 0.0
        self.brake_target_engagement_fraction[Position.LEFT] = 1.0
        self.send_signals()

    def brake_right_engage(self):
        self.clutch_target_engagement_fraction[Position.RIGHT] = 0.0
        self.brake_target_engagement_fraction[Position.RIGHT] = 1.0
        self.send_signals()

    def blade_lift_up(self):
        self.blade_lift_target_velocity = -self.blade_lift_velocity_amplitude
        self.send_signals()

    def blade_lift_down(self):
        self.blade_lift_target_velocity = self.blade_lift_velocity_amplitude
        self.send_signals()

    def blade_tilt_up(self):
        self.blade_tilt_target_velocity = -self.blade_tilt_velocity_amplitude
        self.send_signals()

    def blade_tilt_down(self):
        self.blade_tilt_target_velocity = self.blade_tilt_velocity_amplitude
        self.send_signals()

    def ripper_lift_up(self):
        self.ripper_lift_target_velocity = -self.ripper_lift_velocity_amplitude
        self.send_signals()

    def ripper_lift_down(self):
        self.ripper_lift_target_velocity = self.ripper_lift_velocity_amplitude
        self.send_signals()

    def stop_motor(self):
        self.drive_target = 0.0
        if self.gear_input is not None:
            self.target_gear = 0
        self.send_signals()

    def stop_brake(self):
        self.clutch_target_engagement_fraction[Position.LEFT] = 1.0
        self.clutch_target_engagement_fraction[Position.RIGHT] = 1.0
        self.brake_target_engagement_fraction[Position.LEFT] = 0.0
        self.brake_target_engagement_fraction[Position.RIGHT] = 0.0
        self.send_signals()

    def stop_blade_lift(self):
        self.blade_lift_target_velocity = 0.0
        self.send_signals()

    def stop_blade_tilt(self):
        self.blade_tilt_target_velocity = 0.0
        self.send_signals()

    def stop_ripper(self):
        self.ripper_lift_target_velocity = 0.0
        self.send_signals()

    def handle_up(self, key) -> bool:
        if key.down:
            self.forward()
        else:
            self.stop_motor()
        return True

    def handle_down(self, key) -> bool:
        if key.down:
            self.backward()
        else:
            self.stop_motor()
        return True

    def handle_insert(self, key) -> bool:
        if key.down:
            self.brake_left_engage()
        else:
            self.stop_brake()
        return True

    def handle_delete(self, key) -> bool:
        if key.down:
            self.brake_right_engage()
        else:
            self.stop_brake()
        return True

    def handle_pageup(self, key) -> bool:
        if key.down:
            self.blade_lift_up()
        else:
            self.stop_blade_lift()
        return True

    def handle_pagedown(self, key) -> bool:
        if key.down:
            self.blade_lift_down()
        else:
            self.stop_blade_lift()
        return True

    def handle_home(self, key) -> bool:
        if key.down:
            self.blade_tilt_up()
        else:
            self.stop_blade_tilt()
        return True

    def handle_end(self, key) -> bool:
        if key.down:
            self.blade_tilt_down()
        else:
            self.stop_blade_tilt()
        return True

    def handle_right(self, key) -> bool:
        if key.down:
            self.ripper_lift_up()
        else:
            self.stop_ripper()
        return True

    def handle_left(self, key) -> bool:
        if key.down:
            self.ripper_lift_down()
        else:
            self.stop_ripper()
        return True

    def setup_keyboard_listener(self):
        Input.bind(name="up", key=agxSDK.GuiEventListener.KEY_Up, callback=self.handle_up)
        Input.bind(
            name="down",
            key=agxSDK.GuiEventListener.KEY_Down,
            callback=self.handle_down,
        )
        Input.bind(
            name="insert",
            key=agxSDK.GuiEventListener.KEY_Insert,
            callback=self.handle_insert,
        )
        Input.bind(
            name="delete",
            key=agxSDK.GuiEventListener.KEY_Delete,
            callback=self.handle_delete,
        )
        Input.bind(
            name="pageup",
            key=agxSDK.GuiEventListener.KEY_Page_Up,
            callback=self.handle_pageup,
        )
        Input.bind(
            name="pagedown",
            key=agxSDK.GuiEventListener.KEY_Page_Down,
            callback=self.handle_pagedown,
        )
        Input.bind(
            name="home",
            key=agxSDK.GuiEventListener.KEY_Home,
            callback=self.handle_home,
        )
        Input.bind(
            name="end",
            key=agxSDK.GuiEventListener.KEY_End,
            callback=self.handle_end,
        )
        Input.bind(
            name="right",
            key=agxSDK.GuiEventListener.KEY_Right,
            callback=self.handle_right,
        )
        Input.bind(
            name="left",
            key=agxSDK.GuiEventListener.KEY_Left,
            callback=self.handle_left,
        )


def initialize_scene_decorator():
    camera_data = application().getCameraData()
    camera_data.eye = agx.Vec3(-12.8402, 18.6618, 5.0588)
    camera_data.center = agx.Vec3(1.0282, -0.5451, 2.0728)
    camera_data.up = agx.Vec3(0.0, 0.0, 0.9917)
    camera_data.nearClippingPlane = 0.01
    camera_data.farClippingPlane = 500
    application().applyCameraData(camera_data)
    application().getSceneDecorator().setEnableShadows(True)
    application().getSceneDecorator().setShadowMethod(
        agxOSG.SceneDecorator.SOFT_SHADOWMAP
    )
    application().getSceneDecorator().setEnableShaderState(True)

    for light_idx in range(3):
        application().getSceneDecorator().setEnableCalculateLightPositions(
            light_idx, False
        )

    application().getSceneDecorator().getLightSource(
        agxOSG.SceneDecorator.LIGHT0
    ).setPosition(agx.Vec4(10, 10, 10, 1))
    application().getSceneDecorator().getLightSource(
        agxOSG.SceneDecorator.LIGHT1
    ).setPosition(agx.Vec4(-10, 10, 10, 1))
    application().getSceneDecorator().getLightSource(
        agxOSG.SceneDecorator.LIGHT2
    ).setPosition(agx.Vec4(10, 19, 10, 1))


def buildScene():  # pylint: disable=invalid-name
    initialize_scene_decorator()

    file_dir = os.path.dirname(os.path.abspath(__file__))
    load_result = application().loadOpenPlxFile(
        os.path.join(file_dir, "Scenes", "DozerOnPlane.openplx"),
        root(),
        agxOpenPLX.OptParams().withBundlePath(model_bundle_path()),
    )
    openplx_scene = load_result.scene()

    drive_control_mode = DriveControlMode.ENGINE_THROTTLE
    try:
        drive_input = get_required_object(
            openplx_scene,
            "dozer.drive_train.engine.throttle_input",
            "dozer.drive_train.motor.target_angular_velocity_input",
        )
    except AssertionError:
        drive_input = get_required_object(
            openplx_scene,
            "dozer.drive_train.motor.target_angular_velocity_input",
            "dozer.drive_train.engine.throttle_input",
        )

    if drive_input.getName().endswith("target_angular_velocity_input"):
        drive_control_mode = DriveControlMode.MOTOR_VELOCITY

    gear_input = openplx_scene.getObject("dozer.drive_train.gear_box.gear_selection_input")

    clutch_input = {
        Position.LEFT: openplx_scene.getObject(
            "dozer.drive_train.left_clutch.engagement_fraction_input"
        ),
        Position.RIGHT: openplx_scene.getObject(
            "dozer.drive_train.right_clutch.engagement_fraction_input"
        ),
    }
    assert clutch_input[Position.LEFT]
    assert clutch_input[Position.RIGHT]

    brake_input = {
        Position.LEFT: openplx_scene.getObject(
            "dozer.drive_train.left_brake.engagement_fraction_input"
        ),
        Position.RIGHT: openplx_scene.getObject(
            "dozer.drive_train.right_brake.engagement_fraction_input"
        ),
    }
    assert brake_input[Position.LEFT]
    assert brake_input[Position.RIGHT]

    blade_left_lift_input = get_linear_velocity_input(
        openplx_scene,
        "dozer.blade_connection.left_lift_actuator",
        "dozer.LeftLiftPrismatic_motor",
    )
    blade_right_lift_input = get_linear_velocity_input(
        openplx_scene,
        "dozer.blade_connection.right_lift_actuator",
        "dozer.RightLiftPrismatic_motor",
    )
    blade_left_tilt_input = get_linear_velocity_input(
        openplx_scene,
        "dozer.blade_connection.left_tilt_actuator",
        "dozer.LeftTiltPrismatic_motor",
    )
    blade_right_tilt_input = get_linear_velocity_input(
        openplx_scene,
        "dozer.blade_connection.right_tilt_actuator",
        "dozer.RightTiltPrismatic_motor",
    )
    ripper_left_lift_input = get_optional_linear_velocity_input(
        openplx_scene,
        "dozer.ripper_connection.left_lift_actuator",
        "dozer.LeftRipperLiftPrismatic_motor",
    )
    ripper_right_lift_input = get_optional_linear_velocity_input(
        openplx_scene,
        "dozer.ripper_connection.right_lift_actuator",
        "dozer.RightRipperLiftPrismatic_motor",
    )

    bulldozer_maneuver = BulldozerManeuver(
        input_queue=load_result.getInputSignalQueue(),
        drive_input=drive_input,
        drive_control_mode=drive_control_mode,
        gear_input=gear_input,
        brake_control={
            "clutch_left_input": clutch_input[Position.LEFT],
            "clutch_right_input": clutch_input[Position.RIGHT],
            "brake_left_input": brake_input[Position.LEFT],
            "brake_right_input": brake_input[Position.RIGHT],
        },
        blade_control={
            "blade_left_lift_input": blade_left_lift_input,
            "blade_right_lift_input": blade_right_lift_input,
            "blade_left_tilt_input": blade_left_tilt_input,
            "blade_right_tilt_input": blade_right_tilt_input,
        },
        ripper_control={
            "ripper_left_lift_input": ripper_left_lift_input,
            "ripper_right_lift_input": ripper_right_lift_input,
        },
    )
    bulldozer_maneuver.setup_keyboard_listener()

    content = SimulationContent(simulation=simulation())
    for contact_material in content.contactMaterials.values():
        content.printContactMaterial(contact_material)


init = init_app(name="__main__", scenes=[(buildScene, "1")], autoStepping=True)
