from typing import TypeVar, Generic, Callable
from dataclasses import dataclass

from ...core.actor import Actor
from ...core.rollout import Rollout


U = TypeVar("U")
"""The type of the actor trajectory update object inside the rollout."""

R = TypeVar("R")
"""The type of the rollout result."""


@dataclass
class _ResultMixin(Generic[R]):
    result: R


class RolloutActor(Generic[U, R], _ResultMixin[R], Actor[tuple[str, U]]):
    """
    This class wraps a rollout object so that it appears as an actor.
    When the actor finishes running, the final result is gathered and the rollout resources are cleaned up.
    """

    async def initialize(self) -> None:
        pass

    async def run(self, on_trajectory_update: Callable[[tuple[str, U]], None] | None = None) -> None:
        await self._rollout.run(on_trajectory_update=on_trajectory_update)

    async def cleanup(self) -> None:
        pass

    @property
    def rollout(self) -> Rollout[U, R]:
        return self._rollout

    @property
    def result(self) -> R:
        return self._rollout.result

    def __init__(self, rollout: Rollout[U, R]):
        super().__init__()

        self._rollout = rollout
