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

from .actor import Actor

U = TypeVar("U")
"""The type of the trajectory update, shared across loggable actors in a world."""

R = TypeVar("R")


class Rollout(Generic[U, R], ABC):
    """
    The rollout represents an experiment in which multiple actors and act and interact.
    This concept roughly maps to "scene" in game development.
    """

    @abstractmethod
    async def build(self) -> None:
        """
        This is where you create your actors and build the world.
        :return:
        """

        raise NotImplementedError()

    @abstractmethod
    async def collect_result(self) -> None:
        """
        This is where you collect the final results (e.g., code submitted) after all actors finishes running.
        You must store the results into `self.result` attribute.
        :return:
        """

        raise NotImplementedError()

    @abstractmethod
    async def cleanup(self) -> None:
        """
        This is where you do the environment-specific cleanup jobs.
        Actor cleanups are handled automatically.
        :return:
        """

    def __init__(self):
        self.actors: dict[str, Actor[U]] = {}
        self.result: R | None = None

    async def run(self, on_trajectory_update: Callable[[str, U], None] = None) -> None:
        """
        This is an all-in-one method.
        Builds the world, runs the simulation, collects final results and cleans up resources.

        Simulation ends when all loggable actors signal end of trajectory update iterator.
        If there are no loggable actors, then simulation ends immediately.
        :param on_trajectory_update: Called each time a new trajectory update occurs.
        First argument is agent ID; second is trajectory update object.
        :return: The trajectory iterator. Keys are actor IDs; values are trajectory updates.
        """

        # Build the world
        await self.build()

        # Initialize all actors
        await asyncio.gather(*[actor.initialize() for actor in self.actors.values()])

        # Run the actors until completion
        async def run_actor(actor_id: str, actor: Actor):
            def on_update(update_object: U):
                on_trajectory_update(actor_id, update_object)

            await actor.run(on_trajectory_update=on_update)

        await asyncio.gather(*[run_actor(actor_id, actor) for actor_id, actor in self.actors.items()])

        # Collect final results
        await self.collect_result()

        # Cleanup environment
        await self.cleanup()

        # Cleanup actors
        await asyncio.gather(*[actor.cleanup() for actor in self.actors.values()])
