from abc import ABC, abstractmethod
from typing import Generic, TypeVar, Callable

U = TypeVar('U')
"""The type of the rollout trajectory update."""


# TODO: Replace initialize/cleanup lifecycle management with an async context manager like those in FastAPI
class Actor(Generic[U], ABC):
    """
    Abstract class for actors.
    Actor is an abstract concept;
    could be an agent, could be the environment, could be an object in the environment etc.
    This concept roughly maps to "actor" or "GameObject" in game development.
    could be concrete agents containing actor loops;
    could be just clients accessing remote agents.
    The important thing is that it provides a uniform API.

    The actor may produce a "trajectory" as it exists.
    At the abstract level, a trajectory is basically a stream of logs.
    For agents, this trajectory could be the growing context, memory, etc.

    The actor is expected to be used like this:

    1. Call `initialize` to initialize and start the actor;
    2. Call `stream_trajectory` to get an async stream of trajectory updates;
    items will automatically appear as the actor works and extends its trajectory.
    """

    @abstractmethod
    async def initialize(self) -> None:
        """
        This is where you initialize the actor.
        This method exists so that logic that needs to be run
        before the actor can interact with other actors can be put here.
        :return:
        """

        raise NotImplementedError()

    @abstractmethod
    async def run(self, on_trajectory_update: Callable[[U], None] | None = None) -> None:
        """
        Runs the actor until completion.
        This is not "fresh start" if the actor is already initialized from some non-default state;
        in this case it's more like "revive the actor from a given state".

        :param on_trajectory_update: Called each time when there's a trajectory update.
        :return:
        """

        raise NotImplementedError()

    @abstractmethod
    async def cleanup(self) -> None:
        """
        This is where you clean up resources when the actor is finished.
        :return:
        """

        raise NotImplementedError()
