#!/usr/bin/env python3
import rospy
from lasr_rasa.srv import Rasa, RasaResponse
import sys
import json
from rasa.core.agent import Agent
from collections import defaultdict
import asyncio

class RasaWrapper:
    def __init__(self, model_path):
        self.agent = Agent.load(model_path)

    async def __call__(self, req):
        response = RasaResponse()
        resp = await self.agent.parse_message(req.text)
        print(resp)
        entities_dict = defaultdict(list)
        for entity_info in resp["entities"]:
            entity_dict = {k: v for k, v in entity_info.items() if k != "entity"}
            entities_dict[entity_info["entity"]].append(entity_dict)
        resp["entities"] = entities_dict
        response.json_response = json.dumps(resp)
        response.success = True
        return response


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: rosrun lasr_rasa service [model_path]")
    rospy.init_node("lasr_rasa")
    rasa = RasaWrapper(sys.argv[1])
    rospy.loginfo("Rasa service is ready!")
    rospy.Service("/lasr_rasa/parse", Rasa, lambda req : asyncio.run(rasa(req)))
    rospy.spin()
