from typing import TypeVar, Generic, Callable, Coroutine, Any
import random
import asyncio

from .utils.rollout_actor import RolloutActor
from ..utils.random_extended import split_rngs
from ..core.rollout import Rollout


U = TypeVar("U")
R = TypeVar("R")


class MultiRunRollout(Generic[U, R], Rollout[tuple[str, U], list[R]]):
    """
    A special rollout that runs multiple sample rollouts concurrently.
    """

    def __init__(
            self,
            rng: random.Random,
            n_rollouts: int,
            build_rollout: Callable[[random.Random], Coroutine[Any, Any, Rollout[U, R]]],
    ):
        """
        Constructor.
        :param build_rollout: Takes a random number generator and builds a rollout.
        """

        super().__init__()

        self.actors: dict[str, RolloutActor[U, R]] = {}
        self._rngs = split_rngs(rng=rng, n_splits=n_rollouts)
        self._build_rollout = build_rollout

    async def build(self) -> None:
        # Build the rollouts
        rollouts: tuple[Rollout[U, R], ...] = await asyncio.gather(*[self._build_rollout(rng) for rng in self._rngs])
        actors = [RolloutActor(rollout=rollout) for rollout in rollouts]

        # Assign the actors
        for i, actor in enumerate(actors):
            self.actors[str(i)] = actor

    async def collect_result(self) -> None:
        self.result = {key: actor.result for key, actor in self.actors.items()}

    async def cleanup(self) -> None:
        pass