from collections.abc import Callable
from typing import ParamSpec, TypeVar, overload

import remedapy as R
from remedapy.decorator import make_data_last

T = TypeVar('T')
ReturnType = TypeVar('ReturnType')
ReturnType2 = TypeVar('ReturnType2')
P = ParamSpec('P')
Predicate = Callable[[T], bool] | Callable[[], bool]
OnTrue = Callable[[T], ReturnType] | Callable[[], ReturnType]


@overload
def when(
    data: T,
    predicate: Predicate[T],
    on_true: OnTrue[T, ReturnType],
    /,
) -> ReturnType | T: ...


@overload
def when(
    predicate: Predicate[T],
    on_true: OnTrue[T, ReturnType],
    /,
) -> Callable[[T], ReturnType | T]: ...


@overload
def when(
    data: T,
    predicate: Predicate[T],
    on_true: OnTrue[T, ReturnType],
    /,
    *,
    on_false: OnTrue[T, ReturnType2],
) -> ReturnType2 | ReturnType | T: ...


@overload
def when(
    predicate: Predicate[T],
    on_true: OnTrue[T, ReturnType],
    /,
    *,
    on_false: OnTrue[T, ReturnType2],
) -> Callable[[T], ReturnType2 | ReturnType | T]: ...


@make_data_last
def when(
    data: T,
    predicate: Callable[..., bool],
    transformer: OnTrue[T, ReturnType],
    /,
    *,
    on_false: OnTrue[T, ReturnType2] | None = None,
) -> ReturnType | ReturnType2 | T:
    """
    Given data, predicate, and transformer; applies the transformer to the data if the predicate is true on the data.

    Accepts optional on_false transformer.

    Predicate and transformers can have arity 0 or 1.

    Parameters
    ----------
    data: T
        Value (positional-only).
    predicate: Callable[[T], bool] | Callable[[], bool]
        The predicate to apply to data (positional-only).
    transformer: Callable[[T], ReturnType] | Callable[[], ReturnType]
        The transformer to apply to data if the predicate is true (positional-only).
    on_false: Callable[[T], ReturnType] | Callable[[], ReturnType] | None
        The transformer to apply to data if the predicate is false (keyword-only, optional).

    Returns
    -------
    ReturnType | T
        The result of the transformer if the predicate is true, otherwise the original data.

    Examples
    --------
    Data first:
    >>> R.when(4, R.gt(3), R.add(1))
    5
    >>> R.when(2, R.gt(3), R.add(1))
    2
    >>> R.when(2, R.gt(3), R.add(1), on_false=R.multiply(2))
    4
    >>> R.when(2, R.gt(3), R.add(1), on_false=R.constant(5))
    5

    Data last:
    >>> R.when(R.gt(3), R.add(1))(4)
    5
    >>> R.when(R.gt(3), R.add(1))(2)
    2
    >>> R.when(R.gt(3), R.add(1), on_false=R.multiply(2))(2)
    4
    >>> R.when(R.gt(3), R.add(1), on_false=R.constant(5))(2)
    5

    """
    apply_to_data = R.apply(data)
    is_true = apply_to_data(predicate)
    if is_true:
        return apply_to_data(transformer)
    if on_false is None:
        return data
    return apply_to_data(on_false)
