# This file is part of Hypothesis, which may be found at
# https://github.com/HypothesisWorks/hypothesis/
#
# Copyright the Hypothesis Authors.
# Individual contributors are listed in AUTHORS.rst and the git log.
#
# This Source Code Form is subject to the terms of the Mozilla Public License,
# v. 2.0. If a copy of the MPL was not distributed with this file, You can
# obtain one at https://mozilla.org/MPL/2.0/.

import copy
import math
from collections.abc import Iterable
from typing import Any, Callable, Optional, Union, overload

from hypothesis import strategies as st
from hypothesis.errors import InvalidArgument
from hypothesis.internal.conjecture import utils as cu
from hypothesis.internal.conjecture.data import ConjectureData
from hypothesis.internal.conjecture.engine import BUFFER_SIZE
from hypothesis.internal.conjecture.junkdrawer import LazySequenceCopy
from hypothesis.internal.conjecture.utils import combine_labels
from hypothesis.internal.filtering import get_integer_predicate_bounds
from hypothesis.internal.reflection import is_identity_function
from hypothesis.strategies._internal.strategies import (
    T3,
    T4,
    T5,
    Ex,
    RecurT,
    SampledFromStrategy,
    SearchStrategy,
    T,
    check_strategy,
    filter_not_satisfied,
)
from hypothesis.strategies._internal.utils import cacheable, defines_strategy
from hypothesis.utils.conventions import UniqueIdentifier


class TupleStrategy(SearchStrategy[tuple[Ex, ...]]):
    """A strategy responsible for fixed length tuples based on heterogeneous
    strategies for each of their elements."""

    def __init__(self, strategies: Iterable[SearchStrategy[Any]]):
        super().__init__()
        self.element_strategies = tuple(strategies)

    def do_validate(self) -> None:
        for s in self.element_strategies:
            s.validate()

    def calc_label(self) -> int:
        return combine_labels(
            self.class_label, *(s.label for s in self.element_strategies)
        )

    def __repr__(self) -> str:
        tuple_string = ", ".join(map(repr, self.element_strategies))
        return f"TupleStrategy(({tuple_string}))"

    def calc_has_reusable_values(self, recur: RecurT) -> bool:
        return all(recur(e) for e in self.element_strategies)

    def do_draw(self, data: ConjectureData) -> tuple[Ex, ...]:
        return tuple(data.draw(e) for e in self.element_strategies)

    def calc_is_empty(self, recur: RecurT) -> bool:
        return any(recur(e) for e in self.element_strategies)


@overload
def tuples() -> SearchStrategy[tuple[()]]:  # pragma: no cover
    ...


@overload
def tuples(__a1: SearchStrategy[Ex]) -> SearchStrategy[tuple[Ex]]:  # pragma: no cover
    ...


@overload
def tuples(
    __a1: SearchStrategy[Ex], __a2: SearchStrategy[T]
) -> SearchStrategy[tuple[Ex, T]]:  # pragma: no cover
    ...


@overload
def tuples(
    __a1: SearchStrategy[Ex], __a2: SearchStrategy[T], __a3: SearchStrategy[T3]
) -> SearchStrategy[tuple[Ex, T, T3]]:  # pragma: no cover
    ...


@overload
def tuples(
    __a1: SearchStrategy[Ex],
    __a2: SearchStrategy[T],
    __a3: SearchStrategy[T3],
    __a4: SearchStrategy[T4],
) -> SearchStrategy[tuple[Ex, T, T3, T4]]:  # pragma: no cover
    ...


@overload
def tuples(
    __a1: SearchStrategy[Ex],
    __a2: SearchStrategy[T],
    __a3: SearchStrategy[T3],
    __a4: SearchStrategy[T4],
    __a5: SearchStrategy[T5],
) -> SearchStrategy[tuple[Ex, T, T3, T4, T5]]:  # pragma: no cover
    ...


@overload
def tuples(
    *args: SearchStrategy[Any],
) -> SearchStrategy[tuple[Any, ...]]:  # pragma: no cover
    ...


@cacheable
@defines_strategy()
def tuples(*args: SearchStrategy[Any]) -> SearchStrategy[tuple[Any, ...]]:
    """Return a strategy which generates a tuple of the same length as args by
    generating the value at index i from args[i].

    e.g. tuples(integers(), integers()) would generate a tuple of length
    two with both values an integer.

    Examples from this strategy shrink by shrinking their component parts.
    """
    for arg in args:
        check_strategy(arg)

    return TupleStrategy(args)


class ListStrategy(SearchStrategy[list[Ex]]):
    """A strategy for lists which takes a strategy for its elements and the
    allowed lengths, and generates lists with the correct size and contents."""

    _nonempty_filters: tuple[Callable[[Any], Any], ...] = (bool, len, tuple, list)

    def __init__(
        self,
        elements: SearchStrategy[Ex],
        min_size: int = 0,
        max_size: Optional[Union[float, int]] = math.inf,
    ):
        super().__init__()
        self.min_size = min_size or 0
        self.max_size = max_size if max_size is not None else math.inf
        assert 0 <= self.min_size <= self.max_size
        self.average_size = min(
            max(self.min_size * 2, self.min_size + 5),
            0.5 * (self.min_size + self.max_size),
        )
        self.element_strategy = elements
        if min_size > BUFFER_SIZE:
            raise InvalidArgument(
                f"{self!r} can never generate an example, because min_size is larger "
                "than Hypothesis supports.  Including it is at best slowing down your "
                "tests for no benefit; at worst making them fail (maybe flakily) with "
                "a HealthCheck error."
            )

    def calc_label(self) -> int:
        return combine_labels(self.class_label, self.element_strategy.label)

    def do_validate(self) -> None:
        self.element_strategy.validate()
        if self.is_empty:
            raise InvalidArgument(
                "Cannot create non-empty lists with elements drawn from "
                f"strategy {self.element_strategy!r} because it has no values."
            )
        if self.element_strategy.is_empty and 0 < self.max_size < float("inf"):
            raise InvalidArgument(
                f"Cannot create a collection of max_size={self.max_size!r}, "
                "because no elements can be drawn from the element strategy "
                f"{self.element_strategy!r}"
            )

    def calc_is_empty(self, recur: RecurT) -> bool:
        if self.min_size == 0:
            return False
        else:
            return recur(self.element_strategy)

    def do_draw(self, data: ConjectureData) -> list[Ex]:
        if self.element_strategy.is_empty:
            assert self.min_size == 0
            return []

        elements = cu.many(
            data,
            min_size=self.min_size,
            max_size=self.max_size,
            average_size=self.average_size,
        )
        result = []
        while elements.more():
            result.append(data.draw(self.element_strategy))
        return result

    def __repr__(self) -> str:
        return (
            f"{self.__class__.__name__}({self.element_strategy!r}, "
            f"min_size={self.min_size:_}, max_size={self.max_size:_})"
        )

    def filter(self, condition: Callable[[list[Ex]], Any]) -> SearchStrategy[list[Ex]]:
        if condition in self._nonempty_filters or is_identity_function(condition):
            assert self.max_size >= 1, "Always-empty is special cased in st.lists()"
            if self.min_size >= 1:
                return self
            new = copy.copy(self)
            new.min_size = 1
            return new

        constraints, pred = get_integer_predicate_bounds(condition)
        if constraints.get("len") and (
            "min_value" in constraints or "max_value" in constraints
        ):
            new = copy.copy(self)
            new.min_size = max(
                self.min_size, constraints.get("min_value", self.min_size)
            )
            new.max_size = min(
                self.max_size, constraints.get("max_value", self.max_size)
            )
            # Unsatisfiable filters are easiest to understand without rewriting.
            if new.min_size > new.max_size:
                return SearchStrategy.filter(self, condition)
            # Recompute average size; this is cheaper than making it into a property.
            new.average_size = min(
                max(new.min_size * 2, new.min_size + 5),
                0.5 * (new.min_size + new.max_size),
            )
            if pred is None:
                return new
            return SearchStrategy.filter(new, condition)

        return SearchStrategy.filter(self, condition)


class UniqueListStrategy(ListStrategy[Ex]):
    def __init__(
        self,
        elements: SearchStrategy[Ex],
        min_size: int,
        max_size: Optional[Union[float, int]],
        # TODO: keys are guaranteed to be Hashable, not just Any, but this makes
        # other things harder to type
        keys: tuple[Callable[[Ex], Any], ...],
        tuple_suffixes: Optional[SearchStrategy[tuple[Ex, ...]]],
    ):
        super().__init__(elements, min_size, max_size)
        self.keys = keys
        self.tuple_suffixes = tuple_suffixes

    def do_draw(self, data: ConjectureData) -> list[Ex]:
        if self.element_strategy.is_empty:
            assert self.min_size == 0
            return []

        elements = cu.many(
            data,
            min_size=self.min_size,
            max_size=self.max_size,
            average_size=self.average_size,
        )
        seen_sets: tuple[set[Ex], ...] = tuple(set() for _ in self.keys)
        # actually list[Ex], but if self.tuple_suffixes is present then Ex is a
        # tuple[T, ...] because self.element_strategy is a TuplesStrategy, and
        # appending a concrete tuple to `result: list[Ex]` makes mypy unhappy
        # without knowing that Ex = tuple.
        result: list[Any] = []

        # We construct a filtered strategy here rather than using a check-and-reject
        # approach because some strategies have special logic for generation under a
        # filter, and FilteredStrategy can consolidate multiple filters.
        def not_yet_in_unique_list(val: Ex) -> bool:  # type: ignore # covariant type param
            return all(key(val) not in seen for key, seen in zip(self.keys, seen_sets))

        filtered = self.element_strategy._filter_for_filtered_draw(
            not_yet_in_unique_list
        )
        while elements.more():
            value = filtered.do_filtered_draw(data)
            if value is filter_not_satisfied:
                elements.reject(f"Aborted test because unable to satisfy {filtered!r}")
            else:
                assert not isinstance(value, UniqueIdentifier)
                for key, seen in zip(self.keys, seen_sets):
                    seen.add(key(value))
                if self.tuple_suffixes is not None:
                    value = (value, *data.draw(self.tuple_suffixes))  # type: ignore
                result.append(value)
        assert self.max_size >= len(result) >= self.min_size
        return result


class UniqueSampledListStrategy(UniqueListStrategy):
    def do_draw(self, data: ConjectureData) -> list[Ex]:
        assert isinstance(self.element_strategy, SampledFromStrategy)

        should_draw = cu.many(
            data,
            min_size=self.min_size,
            max_size=self.max_size,
            average_size=self.average_size,
        )
        seen_sets: tuple[set[Ex], ...] = tuple(set() for _ in self.keys)
        result: list[Any] = []

        remaining = LazySequenceCopy(self.element_strategy.elements)

        while remaining and should_draw.more():
            j = data.draw_integer(0, len(remaining) - 1)
            value = self.element_strategy._transform(remaining.pop(j))
            if value is not filter_not_satisfied and all(
                key(value) not in seen for key, seen in zip(self.keys, seen_sets)
            ):
                for key, seen in zip(self.keys, seen_sets):
                    seen.add(key(value))
                if self.tuple_suffixes is not None:
                    value = (value, *data.draw(self.tuple_suffixes))
                result.append(value)
            else:
                should_draw.reject(
                    "UniqueSampledListStrategy filter not satisfied or value already seen"
                )
        assert self.max_size >= len(result) >= self.min_size
        return result


class FixedDictStrategy(SearchStrategy[dict[Any, Any]]):
    """A strategy which produces dicts with a fixed set of keys, given a
    strategy for each of their equivalent values.

    e.g. {'foo' : some_int_strategy} would generate dicts with the single
    key 'foo' mapping to some integer.
    """

    def __init__(
        self,
        mapping: dict[Any, SearchStrategy[Any]],
        *,
        optional: Optional[dict[Any, SearchStrategy[Any]]],
    ):
        dict_type = type(mapping)
        self.mapping = mapping
        keys = tuple(mapping.keys())
        self.fixed = st.tuples(*[mapping[k] for k in keys]).map(
            lambda value: dict_type(zip(keys, value))
        )
        self.optional = optional

    def do_draw(self, data: ConjectureData) -> dict[Any, Any]:
        value = data.draw(self.fixed)
        if self.optional is None:
            return value

        remaining = [k for k, v in self.optional.items() if not v.is_empty]
        should_draw = cu.many(
            data, min_size=0, max_size=len(remaining), average_size=len(remaining) / 2
        )
        while should_draw.more():
            j = data.draw_integer(0, len(remaining) - 1)
            remaining[-1], remaining[j] = remaining[j], remaining[-1]
            key = remaining.pop()
            value[key] = data.draw(self.optional[key])
        return value

    def calc_is_empty(self, recur: RecurT) -> bool:
        return recur(self.fixed)

    def __repr__(self) -> str:
        if self.optional is not None:
            return f"fixed_dictionaries({self.mapping!r}, optional={self.optional!r})"
        return f"fixed_dictionaries({self.mapping!r})"
