import random
from typing import Dict, Optional, Tuple

from gymnasium.core import ObsType
from pydantic import BaseModel

from primaite.game.agent.actions import ActionManager
from primaite.game.agent.interface import AbstractScriptedAgent
from primaite.game.agent.observations.observation_manager import ObservationManager
from primaite.game.agent.rewards import RewardFunction


class RandomAgent(AbstractScriptedAgent):
    """Agent that ignores its observation and acts completely at random."""

    def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
        """Sample the action space randomly.

        :param obs: Current observation for this agent, not used in RandomAgent
        :type obs: ObsType
        :param timestep: The current simulation timestep, not used in RandomAgent
        :type timestep: int
        :return: Action formatted in CAOS format
        :rtype: Tuple[str, Dict]
        """
        return self.action_manager.get_action(self.action_manager.space.sample())


class PeriodicAgent(AbstractScriptedAgent):
    """Agent that does nothing most of the time, but executes application at regular intervals (with variance)."""

    class Settings(BaseModel):
        """Configuration values for when an agent starts performing actions."""

        start_step: int = 20
        "The timestep at which an agent begins performing it's actions."
        start_variance: int = 5
        "Deviation around the start step."
        frequency: int = 5
        "The number of timesteps to wait between performing actions."
        variance: int = 0
        "The amount the frequency can randomly change to."
        max_executions: int = 999999
        "Maximum number of times the agent can execute its action."

    def __init__(
        self,
        agent_name: str,
        action_space: ActionManager,
        observation_space: ObservationManager,
        reward_function: RewardFunction,
        settings: Optional[Settings] = None,
    ) -> None:
        """Initialise PeriodicAgent."""
        super().__init__(
            agent_name=agent_name,
            action_space=action_space,
            observation_space=observation_space,
            reward_function=reward_function,
        )
        self.settings = settings or PeriodicAgent.Settings()
        self._set_next_execution_timestep(timestep=self.settings.start_step, variance=self.settings.start_variance)
        self.num_executions = 0

    def _set_next_execution_timestep(self, timestep: int, variance: int) -> None:
        """Set the next execution timestep with a configured random variance.

        :param timestep: The timestep when the next execute action should be taken.
        :type timestep: int
        :param variance: Uniform random variance applied to the timestep
        :type variance: int
        """
        random_increment = random.randint(-variance, variance)
        self.next_execution_timestep = timestep + random_increment

    def get_action(self, obs: ObsType, timestep: int) -> Tuple[str, Dict]:
        """Do nothing, unless the current timestep is the next execution timestep, in which case do the action."""
        if timestep == self.next_execution_timestep and self.num_executions < self.settings.max_executions:
            self.num_executions += 1
            self._set_next_execution_timestep(timestep + self.settings.frequency, self.settings.variance)
            return "NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}

        return "DONOTHING", {}