from time import perf_counter
from queue import Empty
from typing import List, Optional

from rlbot.agents.base_agent import BaseAgent
from rlbot.matchcomms.client import MatchcommsClient

from tmcp import TMCP_VERSION
from .message import TMCPMessage


MAX_PACKETS_PER_TICK: int = 50
TIME_BETWEEN_MESSAGES: float = 0.1


class TMCPHandler:
    """The class for handling TMCP.

    Create an instance by just passing your agent in:

    ```
    def initialize_agent(self):
        self.tmcp_handler = TMCPHandler(self)
    ```

    Usage is also very straightforward:

    ```
    def get_output(self, packet: GameTickPacket) -> SimpleControllerState:
        # Receive and parse all new matchcomms messages into TMCPMessage objects.
        new_messages: List[TMCPMessage] = self.tmcp_handler.recv()
        # Handle TMCPMessages.
        for message in new_messages:
            if message.action_type == ActionType.BALL:
                print(message.time)
        
        ...

        # You can send messages like this.
        self.tmcp_handler.send_boost_action(pad_index)

        # Or you can create them and send them more directly:
        my_message = TMCPMessage.ball_action(self.team, self.index, estimated_time_of_arrival)
        self.tmcp_handler.send(my_message)
    ```
    """

    def __init__(self, agent: BaseAgent):
        self.matchcomms: MatchcommsClient = agent.matchcomms
        self.index: int = agent.index
        self.team: int = agent.team
        self.last_time: float = 0.0

    def send(self, message: TMCPMessage) -> bool:
        """Send a TMCPMessage over match comms. Will not send messages if they are coming too quickly.
        Returns whether a message was sent."""
        current_time = perf_counter()
        if current_time - self.last_time < TIME_BETWEEN_MESSAGES:
            return False
        self.matchcomms.outgoing_broadcast.put_nowait(message.to_dict())
        self.last_time: float = current_time
        return True

    def recv(self) -> List[TMCPMessage]:
        messages = []
        # Receive messages until we reach the maximum packets per tick or the queue is empty.
        for _ in range(MAX_PACKETS_PER_TICK):
            try:
                message = self.parse(self.matchcomms.incoming_broadcast.get_nowait())
                if message is not None:
                    messages.append(message)
            except Empty:
                break
        return messages

    def parse(self, message: dict) -> Optional[TMCPMessage]:
        # Ignore messages using a different version of the protocol.
        if message.get("tmcp_version") != TMCP_VERSION:
            return None
        # Ignore messages by opposing team.
        if message.get("team") != self.team:
            return None

        return TMCPMessage.from_dict(message)

    def send_ball_action(self, time: Optional[float] = None) -> bool:
        """The bot is going for the ball.

        `time` - Game time that your bot will arrive at the ball.
        """
        if time is None:
            msg = TMCPMessage.ball_action(self.team, self.index)
        else:
            msg = TMCPMessage.ball_action(self.team, self.index, time)
        return self.send(msg)

    def send_boost_action(self, target: int) -> bool:
        """The bot is going for boost.

        `target` - Index of the boost pad the bot is going to collect.
        """
        return self.send(TMCPMessage.boost_action(self.team, self.index, target))

    def send_demo_action(self, target: int, time: Optional[float] = None) -> bool:
        """The bot is going to demolish another car.

        `target` - Index of the bot that will be demoed.
        `time` - Game time that the bot will demo the other bot.
        """
        if time is None:
            msg = TMCPMessage.demo_action(self.team, self.index, target)
        else:
            msg = TMCPMessage.demo_action(self.team, self.index, target, time)
        return self.send(msg)

    def send_wait_action(self) -> bool:
        """The bot is waiting for a chance to go for the ball or boost.
        Some examples are positioning (retreating/shadowing) and recovering.
        """
        return self.send(TMCPMessage.wait_action(self.team, self.index))
