from __future__ import annotations

from abc import ABC
from typing import Union

import pydantic.dataclasses as pydc

from .transform_spec import EventEncoder
from .transform_spec import LabelEncoder
from .transform_spec import TransformSpec
from tempora.utils.dataclass import pydc_config
from tempora.utils.time import Length
from tempora.utils.time import TimeDelta


TargetTransformSpec = Union[TransformSpec, LabelEncoder, EventEncoder]


class TargetSpec(ABC):
    """Abstract base class for all target specification types."""
    columns: list[str] | None
    transform_spec: TargetTransformSpec | None

    def __post_init__(self) -> None:
        if isinstance(self.transform_spec, (LabelEncoder, EventEncoder)):
            if self.columns is None or len(self.columns) != 1:
                raise ValueError(
                    '"columns" must contain exactly one column when '
                    '"transform_spec" is a LabelEncoder or EventEncoder',
                )


@pydc.dataclass(frozen=True, config=pydc_config)
class PredictionWindowTargetSpec(TargetSpec):
    """Target specification for prediction / lookahead windows.

    Args:
        window_len: Length of the prediction window.
        window_shift: Optional shift between the end of the context window and the
            start of the prediction window.
        columns: Optional subset of target columns to generate.
        transform_spec: Optional transform specification for target-window data.
    """
    window_len: TimeDelta | Length
    window_shift: TimeDelta | Length | None = None
    columns: list[str] | None = None
    transform_spec: TargetTransformSpec | None = None


@pydc.dataclass(frozen=True, config=pydc_config)
class NextStepAheadTargetSpec(TargetSpec):
    """Target specification for next-step-ahead prediction.

    Requires the sampler context window to be defined in rows.

    Args:
        columns: Optional subset of target columns to generate.
        transform_spec: Optional transform specification for target-window data.
    """
    columns: list[str] | None = None
    transform_spec: TargetTransformSpec | None = None


@pydc.dataclass(frozen=True, config=pydc_config)
class ContextWindowTargetSpec(TargetSpec):
    """Target specification based on the current context window.

    Args:
        window_len: Optional target window length. If `None`, defaults to the
            sampler context length.
        window_shift: Optional shift between the end of the context window and the end
            of the target window.
        columns: Optional subset of target columns to generate.
        transform_spec: Optional transform specification for target-window data.
    """
    window_len: TimeDelta | Length | None = None
    window_shift: TimeDelta | Length | None = None
    columns: list[str] | None = None
    transform_spec: TargetTransformSpec | None = None


@pydc.dataclass(frozen=True, config=pydc_config)
class SeriesTargetSpec(TargetSpec):
    """Target specification for a SeriesSampler.

    Args:
        columns: Optional subset of target columns to generate.
        transform_spec: Optional transform specification for target-window data.
    """
    columns: list[str] | None = None
    transform_spec: TargetTransformSpec | None = None


@pydc.dataclass(frozen=True, config=pydc_config)
class TimeToEventTargetSpec(TargetSpec):
    """Target specification for time-to-event (TTE) prediction.

    Args:
        window_len: Length of the event-detection window.
        window_shift: Optional shift between the end of the context window and the
            start of the target window.
        columns: Optional subset of target columns to generate.
        transform_spec: Optional transform specification for target-window data.
        censored: Whether missing events should be treated as censored observations.
    """
    window_len: TimeDelta | Length
    window_shift: TimeDelta | Length | None = None
    columns: list[str] | None = None
    transform_spec: TargetTransformSpec | None = None
    censored: bool = False
