from __future__ import annotations

import base64
import dataclasses as dc
import hashlib
import sys
from typing import Any
from typing import Callable
from typing import get_args
from typing import Literal

from .time import Length
from .time import TimeDelta


WeightingType = Literal['duration', 'inverse_duration', 'num_rows', 'inverse_num_rows']
SERIALIZER_NAME = 'cloudpickle'
SERIALIZATION_VERSION = 1


@dc.dataclass
class SamplerOptions:
    """Options for batch samplers.

    Args:
        use_table_cache: Use the previously cached sampling table for the dataset,
            otherwise compute and cache a new sampling table prior to sampling.
        incremental_table_update: Build the sampling table incrementally during
            sampling, rather than computing the table upfront. This option is only
            applicable to datasets partitioned on "entity_keys".
        allow_partial_segments: Allow sampling of segments whose length is less than the
            sampler context length, e.g. at the boundaries of a time series.
        allow_null_entity_keys: Allow sampling of segments where one or more entity key
            columns are null.
        weight_series: Optional series-level sample weighting. Use 'duration' / 'num_rows'
            to sample in proportion to series duration (longer series yield more segments),
            'inverse_duration' / 'inverse_num_rows' to favor shorter series, or None (default)
            for approximately equal contribution per series.
        left_censor_len: Exclude the first left_censor_len of each time series from
            sampling.
        right_censor_len: Exclude the last right_censor_len of each time series from
            sampling.
        align_on_data: Align the start of each sampled segment to the nearest time
            point (useful for unevenly sampled data).
        segments_per_query: Number of sampled segments to generate per SQL query.
        segments_buffer_len: Number of sampled segments to buffer from the server.
        max_attempts: Maximum number of attempts to successfully sample a segment from
            the dataset before raising an exception.
        rng_seed: RNG seed for the sampler.
    """  # noqa: E501
    use_table_cache: bool = True
    incremental_table_update: bool = True
    allow_partial_segments: bool = True
    allow_null_entity_keys: bool = False
    weight_series: WeightingType | None = None
    left_censor_len: TimeDelta | Length | None = None
    right_censor_len: TimeDelta | Length | None = None
    align_on_data: bool = False
    segments_per_query: int | None = None
    segments_buffer_len: int | None = None
    max_attempts: int = 1000
    rng_seed: int = 201174

    def __post_init__(self) -> None:
        if self.weight_series and self.weight_series not in get_args(WeightingType):
            raise ValueError(
                f'Invalid value for "weight_series" option: "{self.weight_series}"',
            )

        if self.incremental_table_update and self.weight_series:
            raise ValueError(
                'Sampling table option "incremental_table_update=True" '
                'is not compatible with series weighting',
            )


def serialize_callable(func: Callable[..., Any]) -> dict[str, Any]:
    """Serialize a Python callable to a JSON-serializable payload."""
    try:
        import cloudpickle
    except ImportError as e:
        raise ImportError(
            'cloudpickle is required for callable transform serialization',
        ) from e

    payload = cloudpickle.dumps(func)
    callable_name = getattr(func, '__name__', type(func).__name__)
    callable_qualname = getattr(func, '__qualname__', callable_name)
    callable_module = getattr(func, '__module__', type(func).__module__)

    return {
        'class': 'SerializedCallable',
        'serializer': SERIALIZER_NAME,
        'serialization_version': SERIALIZATION_VERSION,
        'python_version': f'{sys.version_info.major}.{sys.version_info.minor}',
        'callable_name': callable_name,
        'callable_qualname': callable_qualname,
        'callable_module': callable_module,
        'payload_b64': base64.b64encode(payload).decode('ascii'),
        'payload_sha256': hashlib.sha256(payload).hexdigest(),
    }
