import asyncio
import pydantic
import traceback
from typing import Awaitable, cast, Iterable, Iterator, overload

from .types import MessagesAndChoices


MetadataValue = float | int | str | bool | None


class PydanticException(pydantic.BaseModel):
    type: str
    message: str
    traceback: str


class Trajectory(pydantic.BaseModel):
    messages_and_choices: MessagesAndChoices
    reward: float
    metrics: dict[str, float] = {}
    metadata: dict[str, MetadataValue] = {}


class TrajectoryGroup(pydantic.BaseModel):
    trajectories: list[Trajectory]
    metadata: dict[str, MetadataValue] = {}
    exceptions: list[PydanticException] = []

    def __init__(
        self,
        trajectories: Iterable[Trajectory] | Iterable[Awaitable[Trajectory]],
        *,
        metadata: dict[str, MetadataValue] = {},
        exceptions: list[BaseException] = [],
    ) -> None:
        super().__init__(
            trajectories=list(trajectories),
            metadata=metadata,
            exceptions=[
                PydanticException(
                    type=str(type(e)),
                    message=str(e),
                    traceback="\n".join(
                        traceback.format_exception(type(e), e, e.__traceback__)
                    ),
                )
                for e in exceptions
            ],
        )

    def __iter__(self) -> Iterator[Trajectory]:
        return iter(self.trajectories)

    def __len__(self) -> int:
        return len(self.trajectories)

    @overload
    def __new__(
        cls,
        trajectories: Iterable[Trajectory],
        *,
        metadata: dict[str, MetadataValue] = {},
        exceptions: list[BaseException] = [],
    ) -> "TrajectoryGroup": ...

    @overload
    def __new__(
        cls,
        trajectories: Iterable[Awaitable[Trajectory]],
        *,
        metadata: dict[str, MetadataValue] = {},
        exceptions: list[BaseException] = [],
    ) -> Awaitable["TrajectoryGroup"]: ...

    def __new__(
        cls,
        trajectories: Iterable[Trajectory] | Iterable[Awaitable[Trajectory]],
        *,
        metadata: dict[str, MetadataValue] = {},
        exceptions: list[BaseException] = [],
    ) -> "TrajectoryGroup | Awaitable[TrajectoryGroup]":
        ts = list(trajectories)
        if all(isinstance(t, Trajectory) for t in ts):
            group = super().__new__(cls)
            group.__init__(
                trajectories=cast(list[Trajectory], ts),
                metadata=metadata,
                exceptions=exceptions,
            )
            return group
        else:

            async def _(exceptions: list[BaseException]):
                from .gather import get_gather_context, record_metrics

                context = get_gather_context()
                trajectories = []
                for future in asyncio.as_completed(
                    cast(list[Awaitable[Trajectory]], ts)
                ):
                    try:
                        trajectory = await future
                        trajectories.append(trajectory)
                        record_metrics(context, trajectory)
                        context.update_pbar(n=1)
                    except BaseException as e:
                        exceptions.append(e)
                        context.metric_sums["exceptions"] += 1
                        context.update_pbar(n=0)
                        if context.too_many_exceptions():
                            raise
                return TrajectoryGroup(
                    trajectories=trajectories,
                    exceptions=exceptions,
                    metadata=metadata,
                )

            class CoroutineWithMetadata:
                def __init__(self, coro, num_trajectories):
                    self.coro = coro
                    self._num_trajectories = num_trajectories

                def __await__(self):
                    return self.coro.__await__()

            coro = _(exceptions.copy())
            return CoroutineWithMetadata(coro, len(ts))
