import sys
from functools import partial
from typing import Optional, TypeVar, cast

from hypothesis import given, settings
from hypothesis import strategies as st

from hypothesis_awkward.strategies import StMinMaxValuesFactory, none_or, ranges
from hypothesis_awkward.util import safe_compare as sc
from hypothesis_awkward.util import safe_max

if sys.version_info >= (3, 11):
    from typing import Generic, TypedDict
else:
    from typing_extensions import Generic, TypedDict

T = TypeVar('T')


def min_max_starts(
    st_: StMinMaxValuesFactory[T],
) -> st.SearchStrategy[tuple[Optional[T], Optional[T]]]:
    def mins() -> st.SearchStrategy[Optional[T]]:
        return none_or(st_())

    def maxes(min_value: Optional[T]) -> st.SearchStrategy[Optional[T]]:
        return none_or(st_(min_value=min_value))

    return mins().flatmap(lambda min_: st.tuples(st.just(min_), maxes(min_)))


def min_max_ends(
    st_: StMinMaxValuesFactory[T],
    min_start: Optional[T] = None,
) -> st.SearchStrategy[tuple[Optional[T], Optional[T]]]:
    def mins() -> st.SearchStrategy[Optional[T]]:
        return none_or(st_(min_value=min_start))

    def maxes(min_value: Optional[T]) -> st.SearchStrategy[Optional[T]]:
        min_value = safe_max((min_value, min_start))
        return none_or(st_(min_value=min_value))

    return mins().flatmap(lambda min_: st.tuples(st.just(min_), maxes(min_)))


class RangesKwargs(TypedDict, Generic[T], total=False):
    # st_: StMinMaxValuesFactory[T]
    min_start: Optional[T]
    max_start: Optional[T]
    min_end: Optional[T]
    max_end: Optional[T]
    allow_start_none: bool
    allow_end_none: bool
    let_end_none_if_start_none: bool
    allow_equal: bool


@st.composite
def ranges_kwargs(
    draw: st.DrawFn, st_: StMinMaxValuesFactory[T] | None = None
) -> RangesKwargs[T]:
    if st_ is None:
        st_ = st.integers  # type: ignore

    min_start, max_start = draw(min_max_starts(st_=st_))  # type: ignore
    min_end, max_end = draw(min_max_ends(st_=st_, min_start=min_start))  # type: ignore

    drawn = (
        ('min_start', min_start),
        ('max_start', max_start),
        ('min_end', min_end),
        ('max_end', max_end),
    )

    kwargs = draw(
        st.fixed_dictionaries(
            {k: st.just(v) for k, v in drawn if v is not None},
            optional={
                'allow_start_none': st.booleans(),
                'allow_end_none': st.booleans(),
                'let_end_none_if_start_none': st.booleans(),
                'allow_equal': st.booleans(),
            },
        )
    )

    return cast(RangesKwargs[T], kwargs)


st_floats = partial(st.floats, allow_nan=False, allow_infinity=False)


@given(st.data())
def test_ranges_kwargs(data: st.DataObject) -> None:
    st_ = data.draw(st.sampled_from([None, st_floats]))
    kwargs = data.draw(ranges_kwargs(st_=st_))  # type: ignore

    min_start = kwargs.get('min_start')
    max_start = kwargs.get('max_start')
    assert sc(min_start) <= sc(max_start)

    min_end = kwargs.get('min_end')
    max_end = kwargs.get('max_end')
    assert sc(min_start) <= sc(min_end) <= sc(max_end)


@given(st.data())
@settings(max_examples=1000)
def test_ranges(data: st.DataObject) -> None:
    st_ = data.draw(st.sampled_from([None, st_floats]))
    kwargs = data.draw(ranges_kwargs(st_=st_))  # type: ignore

    args = (st_,) if st_ is not None else ()

    start, end = data.draw(ranges(*args, **kwargs))  # type: ignore

    allow_start_none = kwargs.get('allow_start_none', True)
    if not allow_start_none:
        assert start is not None

    let_end_none_if_start_none = kwargs.get('let_end_none_if_start_none', False)
    allow_end_none = kwargs.get('allow_end_none', True)
    if start is None and let_end_none_if_start_none:
        assert end is None
    elif not allow_end_none:
        assert end is not None

    allow_equal = kwargs.get('allow_equal', True)
    if allow_equal:
        assert sc(start) <= sc(end)
    else:
        assert sc(start) < sc(end)

    min_start = kwargs.get('min_start')
    max_start = kwargs.get('max_start')
    assert sc(min_start) <= sc(start) <= sc(max_start)

    min_end = kwargs.get('min_end')
    max_end = kwargs.get('max_end')
    assert sc(min_end) <= sc(end) <= sc(max_end)
