from __future__ import annotations

import dataclasses as dc
from typing import Iterator

import pydantic.dataclasses as pydc

from .base import ClassSamplingSpec
from .base import SamplerBase
from .target_spec import SeriesTargetSpec
from .target_spec import TargetSpec  # noqa: F401
from tempora.datasets.protocols import IsDataset
from tempora.utils.batch import Batch
from tempora.utils.dataclass import pydc_config
from tempora.utils.time import Length
from tempora.utils.time import TimeDelta


@dc.dataclass
class _ContextLenArg:
    context_len: TimeDelta | Length


@pydc.dataclass(config=pydc_config)
class RandomSampler(SamplerBase, _ContextLenArg):
    """Randomly sample batches from a dataset.

    Args:
        context_len: Length of each sampled context window.
    """

    class_sampling: list[ClassSamplingSpec] | None = None

    def __post_init__(self) -> None:
        if self.class_sampling is not None:
            if not self.target_spec:
                raise ValueError(
                    '"class_sampling" requires "target_spec" to be configured',
                )

            class_names = [class_spec.name for class_spec in self.class_sampling]
            if len(class_names) != len(set(class_names)):
                raise ValueError(
                    '"class_sampling" class names must be unique',
                )

            total_weight = sum(class_spec.weight for class_spec in self.class_sampling)
            if total_weight > 1:
                raise ValueError(
                    f'Sum of "class_sampling" weights must be <= 1 (got {total_weight})',  # noqa: E501
                )

        super().__post_init__()


@pydc.dataclass(config=pydc_config)
class SequentialSampler(SamplerBase, _ContextLenArg):
    """Sample batches sequentially from each data series.

    Args:
        random_start: If `True`, randomize the sampling start offset per series.
        random_end: If `True`, randomize the sampling end offset per series.
    """

    random_start: bool = False
    random_end: bool = False

    def __post_init__(self) -> None:
        super().__post_init__()

    def iter_series(self, dataset: IsDataset, /) -> Iterator[Iterator[Batch]]:
        """
        Returns an iterator over batch series generators. Each batch series generator
        generates sequential batches from a specific data series (indexed on a sampled
        entity key value) and stops when all batches for that series have been yielded
        """
        self._dataset = dataset
        self._init_sampler()
        self._start_segment_reader()
        while True:
            # Generator of batches for a specific data series / entity key
            yield self._sample_batches(stop_on_sentinel=True)


@pydc.dataclass(config=pydc_config)
class SeriesSampler(SamplerBase):
    """Sample complete data series as individual batch examples.

    Only `SeriesTargetSpec` is supported for `target_spec` in this sampler.
    """

    target_spec: SeriesTargetSpec | None = None  # type: ignore[assignment]

    def __post_init__(self) -> None:
        super().__post_init__()
