from __future__ import annotations

from collections.abc import Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, overload

import hypothesis.strategies as st
from hypothesis.errors import InvalidArgument

from polars import select, when
from polars._utils.deprecation import issue_deprecation_warning
from polars.dataframe import DataFrame
from polars.datatypes import Array, Boolean, DataType, DataTypeClass, List, Null, Struct
from polars.series import Series
from polars.testing.parametric.strategies._utils import flexhash
from polars.testing.parametric.strategies.data import data
from polars.testing.parametric.strategies.dtype import _instantiate_dtype, dtypes

if TYPE_CHECKING:
    from collections.abc import Collection, Sequence
    from typing import Literal

    from hypothesis.strategies import DrawFn, SearchStrategy

    from polars import LazyFrame
    from polars._typing import PolarsDataType


_ROW_LIMIT = 5  # max generated frame/series length
_COL_LIMIT = 5  # max number of generated cols


@st.composite
def series(
    draw: DrawFn,
    /,
    *,
    name: str | SearchStrategy[str] | None = None,
    dtype: PolarsDataType | None = None,
    min_size: int = 0,
    max_size: int = _ROW_LIMIT,
    strategy: SearchStrategy[Any] | None = None,
    allow_null: bool = True,
    allow_chunks: bool = True,
    allow_masked_out: bool = True,
    unique: bool = False,
    allowed_dtypes: Collection[PolarsDataType] | PolarsDataType | None = None,
    excluded_dtypes: Collection[PolarsDataType] | PolarsDataType | None = None,
    allow_time_zones: bool = True,
    **kwargs: Any,
) -> Series:
    """
    Hypothesis strategy for producing Polars Series.

    .. warning::
        This functionality is currently considered **unstable**. It may be
        changed at any point without it being considered a breaking change.

    Parameters
    ----------
    name : {str, strategy}, optional
        literal string or a strategy for strings (or None), passed to the Series
        constructor name-param.
    dtype : PolarsDataType, optional
        a valid polars DataType for the resulting series.
    min_size : int
        if not passing an exact size, can set a minimum here (defaults to 0).
        no-op if `size` is set.
    max_size : int
        if not passing an exact size, can set a maximum value here (defaults to
        MAX_DATA_SIZE). no-op if `size` is set.
    strategy : strategy, optional
        supports overriding the default strategy for the given dtype.
    allow_null : bool
        Allow nulls as possible values and allow the `Null` data type by default.
    allow_chunks : bool
        Allow the Series to contain multiple chunks.
    allow_masked_out : bool
        Allow the nulls to contain masked out elements.
    unique : bool, optional
        indicate whether Series values should all be distinct.
    allowed_dtypes : {list,set}, optional
        when automatically generating Series data, allow only these dtypes.
    excluded_dtypes : {list,set}, optional
        when automatically generating Series data, exclude these dtypes.
    allow_time_zones
        Allow generating `Datetime` Series with a time zone.
    **kwargs
        Additional keyword arguments that are passed to the underlying data generation
        strategies.

    size : int, optional
        if set, creates a Series of exactly this size (ignoring min_size/max_size
        params).

        .. deprecated:: 1.0.0
            Use `min_size` and `max_size` instead.

    null_probability : float
        Percentage chance (expressed between 0.0 => 1.0) that any Series value is null.
        This is applied independently of any None values generated by the underlying
        strategy.

        .. deprecated:: 0.20.26
            Use `allow_null` instead.

    allow_infinities : bool, optional
        Allow generation of +/-inf values for floating-point dtypes.

        .. deprecated:: 0.20.26
            Use `allow_infinity` instead.

    Notes
    -----
    In actual usage this is deployed as a unit test decorator, providing a strategy
    that generates multiple Series with the given dtype/size characteristics for the
    unit test. While developing a strategy/test, it can also be useful to call
    `.example()` directly on a given strategy to see concrete instances of the
    generated data.

    Examples
    --------
    The strategy is generally used to generate series in a unit test:

    >>> from polars.testing.parametric import series
    >>> from hypothesis import given
    >>> @given(s=series(min_size=3, max_size=5))
    ... def test_series_len(s: pl.Series) -> None:
    ...     assert 3 <= s.len() <= 5

    Drawing examples interactively is also possible with the `.example()` method.
    This should be avoided while running tests.

    >>> from polars.testing.parametric import lists
    >>> s = series(strategy=lists(pl.String, select_from=["xx", "yy", "zz"]))
    >>> s.example()  # doctest: +SKIP
    shape: (4,)
    Series: '' [list[str]]
    [
            ["zz", "zz"]
            ["zz", "xx", "yy"]
            []
            ["xx"]
    ]
    """
    if (null_prob := kwargs.pop("null_probability", None)) is not None:
        allow_null = _handle_null_probability_deprecation(null_prob)  # type: ignore[assignment]
    if (allow_inf := kwargs.pop("allow_infinities", None)) is not None:
        issue_deprecation_warning(
            "`allow_infinities` is deprecated. Use `allow_infinity` instead.",
            version="0.20.26",
        )
        kwargs["allow_infinity"] = allow_inf
    if (chunked := kwargs.pop("chunked", None)) is not None:
        issue_deprecation_warning(
            "`chunked` is deprecated. Use `allow_chunks` instead.",
            version="0.20.26",
        )
        allow_chunks = chunked
    if (size := kwargs.pop("size", None)) is not None:
        issue_deprecation_warning(
            "`size` is deprecated. Use `min_size` and `max_size` instead.",
            version="1.0.0",
        )
        min_size = max_size = size

    if isinstance(allowed_dtypes, (DataType, DataTypeClass)):
        allowed_dtypes = [allowed_dtypes]
    elif allowed_dtypes is not None:
        allowed_dtypes = list(allowed_dtypes)
    if isinstance(excluded_dtypes, (DataType, DataTypeClass)):
        excluded_dtypes = [excluded_dtypes]
    elif excluded_dtypes is not None:
        excluded_dtypes = list(excluded_dtypes)

    if not allow_null and not (allowed_dtypes is not None and Null in allowed_dtypes):
        if excluded_dtypes is None:
            excluded_dtypes = [Null]
        else:
            excluded_dtypes.append(Null)

    if strategy is None:
        if dtype is None:
            dtype_strat = dtypes(
                allowed_dtypes=allowed_dtypes,
                excluded_dtypes=excluded_dtypes,
                allow_time_zones=allow_time_zones,
            )
        else:
            dtype_strat = _instantiate_dtype(
                dtype,
                allowed_dtypes=allowed_dtypes,
                excluded_dtypes=excluded_dtypes,
                allow_time_zones=allow_time_zones,
            )
        dtype = draw(dtype_strat)

    if min_size == max_size:
        size = min_size
    else:
        size = draw(st.integers(min_value=min_size, max_value=max_size))

    if isinstance(name, st.SearchStrategy):
        name = draw(name)

    do_mask_out = (
        allow_masked_out
        and allow_null
        and isinstance(dtype, (List, Array, Struct))
        and draw(st.booleans())
    )

    if size == 0:
        values = []
    else:
        # Create series using dtype-specific strategy to generate values
        if strategy is None:
            strategy = data(
                dtype,  # type: ignore[arg-type]
                allow_null=allow_null and not do_mask_out,
                **kwargs,
            )

        values = draw(
            st.lists(
                strategy,
                min_size=size,
                max_size=size,
                unique_by=(flexhash if unique else None),
            )
        )

    s = Series(name=name, values=values, dtype=dtype)

    # Apply masking out of values
    if do_mask_out:
        values = draw(
            st.lists(
                st.booleans(),
                min_size=size,
                max_size=size,
                unique_by=(flexhash if unique else None),
            )
        )

        mask = Series(name=None, values=values, dtype=Boolean)
        s = select(when(mask).then(s).alias(s.name)).to_series()

    # Apply chunking
    if allow_chunks and size > 1 and draw(st.booleans()):
        split_at = size // 2
        s = s[:split_at].append(s[split_at:])

    return s


@overload
def dataframes(
    cols: int | column | Sequence[column] | None = None,
    *,
    lazy: Literal[False] = ...,
    min_cols: int = 0,
    max_cols: int = _COL_LIMIT,
    min_size: int = 0,
    max_size: int = _ROW_LIMIT,
    include_cols: Sequence[column] | column | None = None,
    allow_null: bool | Mapping[str, bool] = True,
    allow_chunks: bool = True,
    allow_masked_out: bool = True,
    allowed_dtypes: Collection[PolarsDataType] | PolarsDataType | None = None,
    excluded_dtypes: Collection[PolarsDataType] | PolarsDataType | None = None,
    allow_time_zones: bool = True,
    **kwargs: Any,
) -> SearchStrategy[DataFrame]: ...


@overload
def dataframes(
    cols: int | column | Sequence[column] | None = None,
    *,
    lazy: Literal[True],
    min_cols: int = 0,
    max_cols: int = _COL_LIMIT,
    min_size: int = 0,
    max_size: int = _ROW_LIMIT,
    include_cols: Sequence[column] | column | None = None,
    allow_null: bool | Mapping[str, bool] = True,
    allow_chunks: bool = True,
    allow_masked_out: bool = True,
    allowed_dtypes: Collection[PolarsDataType] | PolarsDataType | None = None,
    excluded_dtypes: Collection[PolarsDataType] | PolarsDataType | None = None,
    allow_time_zones: bool = True,
    **kwargs: Any,
) -> SearchStrategy[LazyFrame]: ...


@st.composite
def dataframes(
    draw: DrawFn,
    /,
    cols: int | column | Sequence[column] | None = None,
    *,
    lazy: bool = False,
    min_cols: int = 1,
    max_cols: int = _COL_LIMIT,
    min_size: int = 0,
    max_size: int = _ROW_LIMIT,
    include_cols: Sequence[column] | column | None = None,
    allow_null: bool | Mapping[str, bool] = True,
    allow_chunks: bool = True,
    allow_masked_out: bool = True,
    allowed_dtypes: Collection[PolarsDataType] | PolarsDataType | None = None,
    excluded_dtypes: Collection[PolarsDataType] | PolarsDataType | None = None,
    allow_time_zones: bool = True,
    **kwargs: Any,
) -> DataFrame | LazyFrame:
    """
    Hypothesis strategy for producing Polars DataFrames or LazyFrames.

    .. warning::
        This functionality is currently considered **unstable**. It may be
        changed at any point without it being considered a breaking change.

    Parameters
    ----------
    cols : {int, columns}, optional
        integer number of columns to create, or a sequence of `column` objects
        that describe the desired DataFrame column data.
    lazy : bool, optional
        produce a LazyFrame instead of a DataFrame.
    min_cols : int, optional
        if not passing an exact size, can set a minimum here (defaults to 0).
    max_cols : int, optional
        if not passing an exact size, can set a maximum value here (defaults to
        MAX_COLS).
    min_size : int, optional
        if not passing an exact size, set the minimum number of rows in the
        DataFrame.
    max_size : int, optional
        if not passing an exact size, set the maximum number of rows in the
        DataFrame.
    include_cols : [column], optional
        a list of `column` objects to include in the generated DataFrame. note that
        explicitly provided columns are appended onto the list of existing columns
        (if any present).
    allow_null : bool or Mapping[str, bool]
        Allow nulls as possible values and allow the `Null` data type by default.
        Accepts either a boolean or a mapping of column names to booleans.
    allow_chunks : bool
        Allow the DataFrame to contain multiple chunks.
    allow_masked_out : bool
        Allow the nulls to contain masked out elements.
    allowed_dtypes : {list,set}, optional
        when automatically generating data, allow only these dtypes.
    excluded_dtypes : {list,set}, optional
        when automatically generating data, exclude these dtypes.
    allow_time_zones
        Allow generating `Datetime` columns with a time zone.
    **kwargs
        Additional keyword arguments that are passed to the underlying data generation
        strategies.

    size : int, optional
        if set, will create a DataFrame of exactly this size (and ignore
        the min_size/max_size len params).

        .. deprecated:: 1.0.0
            Use `min_size` and `max_size` instead.

    null_probability : {float, dict[str,float]}, optional
        percentage chance (expressed between 0.0 => 1.0) that a generated value is
        None. this is applied independently of any None values generated by the
        underlying strategy, and can be applied either on a per-column basis (if
        given as a `{col:pct}` dict), or globally. if null_probability is defined
        on a column, it takes precedence over the global value.

        .. deprecated:: 0.20.26
            Use `allow_null` instead.

    allow_infinities : bool, optional
        optionally disallow generation of +/-inf values for floating-point dtypes.

        .. deprecated:: 0.20.26
            Use `allow_infinity` instead.

    Notes
    -----
    In actual usage this is deployed as a unit test decorator, providing a strategy
    that generates DataFrames or LazyFrames with the given characteristics for
    the unit test. While developing a strategy/test, it can also be useful to
    call `.example()` directly on a given strategy to see concrete instances of
    the generated data.

    Examples
    --------
    The strategy is generally used to generate series in a unit test:

    >>> from polars.testing.parametric import dataframes
    >>> from hypothesis import given
    >>> @given(df=dataframes(min_size=3, max_size=5))
    ... def test_df_height(df: pl.DataFrame) -> None:
    ...     assert 3 <= df.height <= 5

    Drawing examples interactively is also possible with the `.example()` method.
    This should be avoided while running tests.

    >>> df = dataframes(allowed_dtypes=[pl.Datetime, pl.Float64], max_cols=3)
    >>> df.example()  # doctest: +SKIP
    shape: (3, 3)
    ┌─────────────┬────────────────────────────┬───────────┐
    │ col0        ┆ col1                       ┆ col2      │
    │ ---         ┆ ---                        ┆ ---       │
    │ f64         ┆ datetime[ns]               ┆ f64       │
    ╞═════════════╪════════════════════════════╪═══════════╡
    │ NaN         ┆ 1844-07-05 06:19:48.848808 ┆ 3.1436e16 │
    │ -1.9914e218 ┆ 2068-12-01 23:05:11.412277 ┆ 2.7415e16 │
    │ 0.5         ┆ 2095-11-19 22:05:17.647961 ┆ -0.5      │
    └─────────────┴────────────────────────────┴───────────┘

    Use :class:`column` for more control over which exactly which columns are generated.

    >>> from polars.testing.parametric import column
    >>> dfs = dataframes(
    ...     [
    ...         column("x", dtype=pl.Int32),
    ...         column("y", dtype=pl.Float64),
    ...     ],
    ...     min_size=2,
    ...     max_size=2,
    ... )
    >>> dfs.example()  # doctest: +SKIP
    shape: (2, 2)
    ┌───────────┬────────────┐
    │ x         ┆ y          │
    │ ---       ┆ ---        │
    │ i32       ┆ f64        │
    ╞═══════════╪════════════╡
    │ -15836    ┆ 1.1755e-38 │
    │ 575050513 ┆ NaN        │
    └───────────┴────────────┘
    """
    if (null_prob := kwargs.pop("null_probability", None)) is not None:
        allow_null = _handle_null_probability_deprecation(null_prob)
    if (allow_inf := kwargs.pop("allow_infinities", None)) is not None:
        issue_deprecation_warning(
            "`allow_infinities` is deprecated. Use `allow_infinity` instead.",
            version="0.20.26",
        )
        kwargs["allow_infinity"] = allow_inf
    if (chunked := kwargs.pop("chunked", None)) is not None:
        issue_deprecation_warning(
            "`chunked` is deprecated. Use `allow_chunks` instead.",
            version="0.20.26",
        )
        allow_chunks = chunked
    if (size := kwargs.pop("size", None)) is not None:
        issue_deprecation_warning(
            "`size` is deprecated. Use `min_size` and `max_size` instead.",
            version="1.0.0",
        )
        min_size = max_size = size

    if isinstance(include_cols, column):
        include_cols = [include_cols]

    if cols is None:
        n_cols = draw(st.integers(min_value=min_cols, max_value=max_cols))
        cols = [column() for _ in range(n_cols)]
    elif isinstance(cols, int):
        cols = [column() for _ in range(cols)]
    elif isinstance(cols, column):
        cols = [cols]
    else:
        cols = list(cols)

    if include_cols:
        cols.extend(list(include_cols))

    if min_size == max_size:
        size = min_size
    else:
        size = draw(st.integers(min_value=min_size, max_value=max_size))

    # Process columns
    for idx, c in enumerate(cols):
        if c.name is None:
            c.name = f"col{idx}"
        if c.allow_null is None:
            if isinstance(allow_null, Mapping):
                c.allow_null = allow_null.get(c.name, True)
            else:
                c.allow_null = allow_null

    allow_series_chunks = draw(st.booleans()) if allow_chunks else False

    data = {
        c.name: draw(
            series(
                name=c.name,
                dtype=c.dtype,
                min_size=size,
                max_size=size,
                strategy=c.strategy,
                allow_null=c.allow_null,  # type: ignore[arg-type]
                allow_chunks=allow_series_chunks,
                allow_masked_out=allow_masked_out,
                unique=c.unique,
                allowed_dtypes=allowed_dtypes,
                excluded_dtypes=excluded_dtypes,
                allow_time_zones=allow_time_zones,
                **kwargs,
            )
        )
        for c in cols
    }

    df = DataFrame(data)

    # Apply chunking
    if allow_chunks and size > 1 and not allow_series_chunks and draw(st.booleans()):
        split_at = size // 2
        df = df[:split_at].vstack(df[split_at:])

    if lazy:
        return df.lazy()

    return df


@dataclass
class column:
    """
    Define a column for use with the `dataframes` strategy.

    .. warning::
        This functionality is currently considered **unstable**. It may be
        changed at any point without it being considered a breaking change.

    Parameters
    ----------
    name : str
        string column name.
    dtype : PolarsDataType
        a polars dtype.
    strategy : strategy, optional
        supports overriding the default strategy for the given dtype.
    allow_null : bool, optional
        Allow nulls as possible values and allow the `Null` data type by default.
    unique : bool, optional
        flag indicating that all values generated for the column should be unique.

    null_probability : float, optional
        percentage chance (expressed between 0.0 => 1.0) that a generated value is
        None. this is applied independently of any None values generated by the
        underlying strategy.

        .. deprecated:: 0.20.26
            Use `allow_null` instead.

    Examples
    --------
    >>> from polars.testing.parametric import column
    >>> dfs = dataframes(
    ...     [
    ...         column("x", dtype=pl.Int32, allow_null=True),
    ...         column("y", dtype=pl.Float64),
    ...     ],
    ...     size=2,
    ... )
    >>> dfs.example()  # doctest: +SKIP
    shape: (2, 2)
    ┌───────────┬────────────┐
    │ x         ┆ y          │
    │ ---       ┆ ---        │
    │ i32       ┆ f64        │
    ╞═══════════╪════════════╡
    │ null      ┆ 1.1755e-38 │
    │ 575050513 ┆ inf        │
    └───────────┴────────────┘
    """

    name: str | None = None
    dtype: PolarsDataType | None = None
    strategy: SearchStrategy[Any] | None = None
    allow_null: bool | None = None
    unique: bool = False

    null_probability: float | None = None

    def __post_init__(self) -> None:
        if self.null_probability is not None:
            self.allow_null = _handle_null_probability_deprecation(  # type: ignore[assignment]
                self.null_probability
            )


def _handle_null_probability_deprecation(
    null_probability: float | Mapping[str, float],
) -> bool | dict[str, bool]:
    issue_deprecation_warning(
        "`null_probability` is deprecated. Use `allow_null` instead.",
        version="0.20.26",
    )

    def prob_to_bool(prob: float) -> bool:
        if not (0.0 <= prob <= 1.0):
            msg = f"`null_probability` should be between 0.0 and 1.0, got {prob!r}"
            raise InvalidArgument(msg)

        return bool(prob)

    if isinstance(null_probability, Mapping):
        return {col: prob_to_bool(prob) for col, prob in null_probability.items()}
    else:
        return prob_to_bool(null_probability)
