from abc import ABC, abstractmethod
from typing import TypedDict, get_type_hints

import pydantic

from silverriver.interfaces import AgentAction, SupportedModes
from silverriver.interfaces.data_models import Observation

# This is to be used as the AgentAction object is being built
ActionDict = TypedDict('ActionDict', get_type_hints(AgentAction))


class AgentState(pydantic.BaseModel):
    pass


class AbstractAgent(ABC):
    """
    Abstract base class for an LLM-based agent.

    This class defines the core interface for an agent powered by a large language model (LLM).
    Subclasses implementing this interface are responsible for determining actions based on
    input observations, post-processing those actions, and managing the agent's internal state.
    """

    @abstractmethod
    def get_action(self, observation: Observation, info: dict, mode: SupportedModes) -> AgentAction:
        """
        Determine the next action for the agent based on the given observation and additional context.
        Optionally specify the LLM model to use for generating the action.

        -----------
        observation : Any
            The current input or observation, which could be a text prompt or other data that the LLM processes.
        info : Any
            Additional context or metadata that might inform the agent's decision-making process.
        mode: Optional[SupportedModels]
            The mode to use for generating the action, which could be a specific LLM model or configuration.
            See SupportedModels for available options

        Returns:
        --------
        AgentAction
            The action selected by the agent, which could be a generated text response or other LLM-driven output.
        """
        raise NotImplementedError

    def get_agent_state(self) -> AgentState:
        """
        Retrieve the current internal state of the agent.

        This state may include aspects like memory, previous interactions, or other data that influence
        the agent's behavior in subsequent operations.

        Returns:
        --------
        AgentState
            The current internal state of the agent, encapsulating memory and other relevant internal data.
        """
        raise NotImplementedError

    def action_postprocess(self, action: ActionDict) -> ActionDict:
        """
        Post-process the action generated by the agent.

        This method is used to refine, log, or otherwise process the action after it has been generated by the agent.

        Parameters:
        -----------
        action : ActionDict
            The action that was generated by the agent, possibly requiring further refinement or handling before use.

        Returns:
        --------
        None
        """
        raise NotImplementedError
