#!/usr/bin/env python3
from bitbots_test.test_case import RosNodeTestCase
from bitbots_test.mocks import MockSubscriber
from actionlib_msgs.msg import GoalStatus
from bitbots_msgs.msg import KickAction, KickGoal, JointCommand
from geometry_msgs.msg import Quaternion
from rclpy.action import ActionClient
import rclpy
from rclpy.node import Node


class KickRunsTestCase(RosNodeTestCase):
    def single_run(self, goal):
        def done_cb(state, result):
            assert state == GoalStatus.SUCCEEDED, "Kick was not successful"
            self.kick_succeeded = True

        self.kick_succeeded = False

        sub = MockSubscriber('kick_motor_goals', JointCommand)
        self.with_assertion_grace_period(lambda: self.assertNodeRunning("dynamic_kick"), 20)
        client = ActionClient(self, KickAction, 'dynamic_kick')
        assert client.wait_for_server(), "Kick action server not running"

        client.send_goal_async(goal)
        client.done_cb = done_cb
        client.wait_for_result()
        sub.wait_until_connected()
        sub.assertMessageReceived()
        assert self.kick_succeeded, "Kick was not successful"

    def test_normal_kick(self):
        goal = KickGoal()
        goal.header.stamp = self.get_clock().now()
        goal.header.frame_id = "base_footprint"
        goal.ball_position.x = 0.2
        goal.kick_direction = Quaternion(0, 0, 0, 1)
        goal.kick_speed = 1.0
        self.single_run(goal)

    def test_unstable_kick(self):
        goal = KickGoal()
        goal.header.stamp = self.get_clock().now()
        goal.header.frame_id = "base_footprint"
        goal.ball_position.x = 0.2
        goal.kick_direction = Quaternion(0, 0, 0, 1)
        goal.kick_speed = 1.0
        goal.unstable = True
        self.single_run(goal)

if __name__ == '__main__':
    from bitbots_test import run_rostests
    run_rostests(KickRunsTestCase)
