from rclpy import Future
from rclpy.node import Node
from rclpy.action import ActionClient
from rclpy.action.client import ClientGoalHandle
from rosbridge_library.internal.message_conversion import extract_values, populate_instance
from btcpp_ros2_interfaces.action import ExecuteTree
from toid_msgs.action import SimpleMoveCoords, SimpleRotate, SimpleTranslateX
from action_msgs.srv import CancelGoal
from tf2_web_republisher_interfaces.action import TFSubscription
from uvicorn.logging import ColourizedFormatter
from fastapi import WebSocket, FastAPI, WebSocketDisconnect
import asyncio
import logging
import threading

log = logging.getLogger("ActionClientHandler")
log.addHandler(logging.StreamHandler())
log.handlers[0].setFormatter(ColourizedFormatter('%(levelprefix)s %(message)s'))
log.setLevel(logging.INFO)
        
class ActionClientHandler:
    action_client: ActionClient
    action_name: str
    running: bool
    connections: set[WebSocket]
    app: FastAPI = None
    lock: threading.Lock

    def __init__(self, app: FastAPI, action_client: ActionClient):
        self.action_client = action_client
        self.action_name = action_client._action_name
        self.app = app
        self.lock = threading.Lock()
        self.running = False
        self.connections = set()

    def send_goal(self, goal):
        with self.lock:
            if not self.action_client.wait_for_server(1):
                log.info("Server doesn't exist yet")
                self.running = False
                return False
            if self.running:
                if not self.cancel_goals():
                    return False
            self.running = True

        log.info("Sending goal")
        future = self.action_client.send_goal_async(
            goal=goal, 
            feedback_callback=self.feedback)
        future.add_done_callback(self.response_callback)
        return True

    def feedback(self, feedback):
        #log.info("Feedback recieved")
        self.sendMessage(
            {
                'type': 'goal_feedback', 
                'name': self.action_name,
                'message': extract_values(feedback.feedback),
            })
        
    def response_callback(self, future: Future):
        log.info("Goal response")
        goal_handle: ClientGoalHandle = future.result()
        if not goal_handle.accepted:
            log.info("Goal not accepeted")
            with self.lock:
                self.running = False
            self.sendMessage(
                {
                    'type': 'goal_accepted', 
                    'name': self.action_name,
                    'accepted': False,
                })
            return

        log.info("Goal accepeted")

        self.sendMessage(
            {
                'type': 'goal_accepted', 
                'name': self.action_name,
                'accepted': True,
            })
        f = goal_handle.get_result_async().add_done_callback(self.done)

    def done(self, future: Future):
        log.info("Goal done maybe")
        result: ExecuteTree.Result = future.result().result

        self.sendMessage(
            {
                'type': 'goal_done', 
                'name': self.action_name,
                'message': extract_values(result),
            })

        with self.lock:
            self.running = False
    
    def sendMessage(self, msg: dict):
        with self.lock:
            connections = list(self.connections)
        for conn in connections:
            try:
                asyncio.run_coroutine_threadsafe(
                    conn.send_json(msg),
                    self.app.state.loop
                )
            except WebSocketDisconnect:
                pass
    
    def cancel_goals(self):
        node: Node = self.action_client._node
        cli = node.create_client(CancelGoal, self.action_name + "/_action/cancel_goal")
        return cli.call(CancelGoal.Request(), 1.0)
        

    def status(self,):
        with self.lock:
            return self.running
    
    def addConnection(self, ws: WebSocket):
        with self.lock:
            self.connections.add(ws)

    def removeConnection(self, ws: WebSocket):
        with self.lock:
            self.connections.remove(ws)
    



actionClient = None

class ServerRunner(Node):
    bt_action_client: ActionClientHandler
    rotate: ActionClientHandler
    move_coords: ActionClientHandler
    move_x: ActionClientHandler
    tf_web: ActionClientHandler
    app: FastAPI = None

    def __init__(self, app: FastAPI):
        super().__init__("RestServerNode")
        self.bt_action_client = ActionClientHandler(app, ActionClient(self, ExecuteTree, "/bt_run"))
        self.rotate = ActionClientHandler(app, ActionClient(self, SimpleRotate, "/rotate"))
        self.move_coords = ActionClientHandler(app, ActionClient(self, SimpleMoveCoords, "/moveCoords"))
        self.move_x = ActionClientHandler(app, ActionClient(self, SimpleTranslateX, "/move_x"))
        self.tf_web = ActionClientHandler(app, ActionClient(self, TFSubscription, "/tf2_web_republisher"))
        self.app = None
        
    def start_tf_pub(self):
        goal = TFSubscription.Goal()
        goal.target_frame = "map"
        goal.source_frames = ["base_footprint"]
        goal.angular_thres = 0.0
        goal.trans_thres = 0.0
        goal.rate = 10.0
        return self.tf_web.send_goal(goal)

    def execute_tree(self, tree_name,):
        goal = ExecuteTree.Goal(target_tree=tree_name)
        return self.bt_action_client.send_goal(goal)

    def rotate_action(self, goal):
        return self.rotate.send_goal(goal)

    def translate_x(self, goal):
        return self.move_x.send_goal(goal)

    def translate_coords(self, goal):
        return self.move_coords.send_goal(goal)

    def add(self, ws: WebSocket):
        self.bt_action_client.addConnection(ws)
        self.rotate.addConnection(ws)
        self.move_coords.addConnection(ws)
        self.move_x.addConnection(ws)
        self.tf_web.addConnection(ws)

    def remove(self, ws: WebSocket):
        self.bt_action_client.removeConnection(ws)
        self.rotate.removeConnection(ws)
        self.move_coords.removeConnection(ws)
        self.tf_web.removeConnection(ws)

    def setRunner(server):
        global actionClient
        actionClient = server
    
    def getRunner():
        global actionClient
        return actionClient
    



    