#!/usr/bin/env python3
import rospy
import argparse
import smach
from lasr_skills.xml_question_answer import XmlQuestionAnswer
from lasr_skills.ask_and_listen import AskAndListen
from lasr_skills.say import Say


class QuestionAnswerStateMachine(smach.StateMachine):
    def __init__(self, input_data: dict):
        smach.StateMachine.__init__(
            self,
            outcomes=["succeeded", "failed"],
            output_keys=["closest_answers"],
        )
        self.userdata.tts_phrase = "I hear you have a question for me; ask away!"
        with self:
            smach.StateMachine.add(
                "GET_QUESTION",
                AskAndListen(),
                transitions={"succeeded": "XML_QUESTION_ANSWER", "failed": "failed"},
                remapping={
                    "tts_phrase:": "tts_phrase",
                    "transcribed_speech": "transcribed_speech",
                },
            )
            smach.StateMachine.add(
                "XML_QUESTION_ANSWER",
                XmlQuestionAnswer(
                    input_data["index_path"],
                    input_data["txt_path"],
                    input_data["xml_path"],
                    input_data["k"],
                ),
                transitions={"succeeded": "SAY_ANSWER", "failed": "failed"},
                remapping={
                    "query_sentence": "transcribed_speech",
                    "closest_answers": "closest_answers",
                },
            )
            smach.StateMachine.add(
                "SAY_ANSWER",
                Say(format_str="The answer to your question is: {}"),
                transitions={
                    "succeeded": "succeeded",
                    "aborted": "failed",
                    "preempted": "failed",
                },
                remapping={"placeholders": "closest_answers"},
            )


def parse_args() -> dict:
    parser = argparse.ArgumentParser(description="GPSR Question Answer")
    parser.add_argument(
        "--k",
        type=int,
        help="The number of closest answers to return",
        required=True,
    )
    parser.add_argument(
        "--index_path",
        type=str,
        help="The path to the index file that is populated with the sentences embeddings of the questions",
        required=True,
    )
    parser.add_argument(
        "--txt_path",
        type=str,
        help="The path to the txt file containing a list of questions.",
        required=True,
    )
    parser.add_argument(
        "--xml_path",
        type=str,
        help="The path to the xml file containing question/answer pairs",
        required=True,
    )
    args, _ = parser.parse_known_args()
    args.k = int(args.k)
    return vars(args)


if __name__ == "__main__":
    rospy.init_node("gpsr_question_answer")
    args: dict = parse_args()
    while not rospy.is_shutdown():
        q_a_sm = QuestionAnswerStateMachine(args)
        outcome = q_a_sm.execute()
        if outcome == "succeeded":
            rospy.loginfo("Question Answer State Machine succeeded")
        else:
            rospy.logerr("Question Answer State Machine failed")

    rospy.spin()
