// FILE: input_generator.rs

use crate::{rand::Rng, report::Observation};

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum NextAttempt<T> {
    Done,
    InfoGathering(T),
    ShrinkAttempt(T),
}

impl<T> NextAttempt<T> {
    pub fn map<U, F: FnOnce(T) -> U>(self, f: F) -> NextAttempt<U> {
        match self {
            Self::Done => NextAttempt::Done,
            Self::InfoGathering(value) => NextAttempt::InfoGathering(f(value)),
            Self::ShrinkAttempt(value) => NextAttempt::ShrinkAttempt(f(value)),
        }
    }
}

pub trait InputGenerator {
    /// The type generated by this generator, and given as input to the code
    /// being tested.
    type Input;

    // TODO(ichen): default this to Self::Input when associated type defaults
    // are stabilized (https://github.com/rust-lang/rust/issues/29661)
    // TODO: explain this
    type InputSource: Clone;

    /// A type that stores shrinking history with the necessary information to
    /// determine the next value to try during shrinking, when to stop
    /// attempting to shrink, and information to generate a report of what
    /// happened during shrinking.
    type History;

    /// Returns the number of possible values of type [`Input`], or [`None`] if
    /// that value is unknown or greater than [`usize::MAX`].
    ///
    /// [`Input`]: NewInputGenerator::Input
    fn cardinality(&self) -> Option<usize>;

    // TODO: docs outdated because ShrinkableInput was added
    /// Returns an iterator over all possible values of type [`Input`], if
    /// possible.
    ///
    /// If [`cardinality`] returns [`None`], the contents of the iterator may be
    /// anything and should not be relied on. If [`cardinality`] returns
    /// [`Some`], the iterator should yield exactly that many items.
    ///
    /// Additionally, assuming [`cardinality`] returns [`Some`], the returned
    /// iterator should yield every possible value of type [`Input`] exactly
    /// once, and always in the same order.
    ///
    /// [`Input`]: NewInputGenerator::Input
    /// [`cardinality`]: NewInputGenerator::cardinality
    fn exhaustive(&self) -> impl Iterator<Item = Self::InputSource>;

    /// Returns the number of adversarial inputs that will be returned by
    /// [`adversarial`], or [`None`] if that value is greater than
    /// [`usize::MAX`] or otherwise should not be relied on.
    ///
    /// [`adversarial`]: NewInputGenerator::adversarial
    fn adversarial_count(&self) -> Option<usize>;

    // TODO: docs outdated because ShrinkableInput was added
    /// Returns an iterator over specifically-chosen adversarial inputs, if
    /// possible. These inputs are ones likely to encounter different branches
    /// or edge cases in the code being tested.
    ///
    /// If [`adversarial_count`] returns [`None`], the contents of the iterator
    /// may be anything and should not be relied on. If [`adversarial_count`]
    /// returns [`Some`], the iterator should yield exactly that many items.
    ///
    /// Additionally, assuming [`adversarial_count`] returns [`Some`], the
    /// returned iterator should each adversarial value exactly once, and always
    /// in the same order.
    ///
    /// [`adversarial_count`]: NewInputGenerator::adversarial_count
    fn adversarial(&self) -> impl Iterator<Item = Self::InputSource>;

    // TODO: docs outdated because ShrinkableInput was added
    /// Samples a random value of type [`Input`].
    ///
    /// The values are not necessarily sampled with equal distribution - some
    /// implementations may choose to return certain values more often than
    /// others, or may even have some values which will never be returned.
    ///
    /// This function should be deterministic when given the same rng in the
    /// same state - that is, if called multiple times with an `rng` that gives
    /// the same outputs, this function should return the same value.
    ///
    /// [`Input`]: NewInputGenerator::Input
    fn sample(&self, rng: &mut (impl Rng + ?Sized)) -> Self::InputSource;

    /// Returns a new [`History`], ready to be used for the shrinking process.
    ///
    /// The `failing_input` argument should be the initial failing input to be
    /// shrunk.
    ///
    /// [`History`]: NewInputGenerator::History
    fn new_history(&self, failing_input: Self::InputSource) -> Self::History;

    /// TODO
    fn current_simplest_failing(&self, history: &Self::History) -> Self::InputSource;

    /// TODO
    fn next_input(
        &self,
        rng: &mut impl Rng,
        history: &Self::History,
    ) -> NextAttempt<Self::InputSource>;

    /// TODO
    fn update_history(
        &self,
        history: &mut Self::History,
        shrinkable_input: Self::InputSource,
        test_passed: bool,
    );

    /// TODO
    fn generate_observations(&self, history: Self::History) -> Vec<Observation>;

    // TODO: docs
    // TODO: should this take an `InputSource` or an `&InputSource`?
    fn create_input(&self, input_source: Self::InputSource) -> Self::Input;
}

// FILE: map.rs
use crate::{input_generator::NextAttempt, InputGenerator};

pub struct Map<G, F> {
    inner_generator: G,
    f: F,
}

impl<U, G: InputGenerator, F: Fn(G::Input) -> U> InputGenerator for Map<G, F> {
    type Input = U;
    type InputSource = G::InputSource;

    type History = G::History;

    fn cardinality(&self) -> Option<usize> {
        self.inner_generator.cardinality()
    }

    fn exhaustive(&self) -> impl Iterator<Item = Self::InputSource> {
        self.inner_generator.exhaustive()
    }

    fn adversarial_count(&self) -> Option<usize> {
        self.inner_generator.adversarial_count()
    }

    fn adversarial(&self) -> impl Iterator<Item = Self::InputSource> {
        self.inner_generator.adversarial()
    }

    fn sample(&self, rng: &mut (impl rand::Rng + ?Sized)) -> Self::InputSource {
        self.inner_generator.sample(rng)
    }

    fn new_history(&self, failing_input: Self::InputSource) -> Self::History {
        self.inner_generator.new_history(failing_input)
    }

    fn current_simplest_failing(&self, history: &Self::History) -> Self::InputSource {
        self.inner_generator.current_simplest_failing(history)
    }

    fn next_input(
        &self,
        rng: &mut impl rand::Rng,
        history: &Self::History,
    ) -> NextAttempt<Self::InputSource> {
        self.inner_generator.next_input(rng, history)
    }

    fn update_history(
        &self,
        history: &mut Self::History,
        shrinkable_input: Self::InputSource,
        test_passed: bool,
    ) {
        self.inner_generator
            .update_history(history, shrinkable_input, test_passed)
    }

    fn generate_observations(&self, history: Self::History) -> Vec<crate::report::Observation> {
        self.inner_generator.generate_observations(history)
    }

    fn create_input(&self, input_source: Self::InputSource) -> Self::Input {
        (self.f)(self.inner_generator.create_input(input_source))
    }
}

impl<U, G: InputGenerator, F: Fn(G::Input) -> U> Map<G, F> {
    pub fn new(inner_generator: G, map_function: F) -> Self {
        Self {
            inner_generator,
            f: map_function,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::run_test;

    #[test]
    fn test_map_does_the_map_thing() {
        let report = run_test(
            |n| n.len() == 1,
            Map::new(
                [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14].as_slice(),
                |t| t.to_string(),
            ),
            &mut crate::rand::thread_rng(),
        )
        .unwrap_err();
        assert_eq!(report.passing_runs, 10);
        assert_eq!(report.shrink_steps, vec![]);
        assert_eq!(report.simplest_failing_input, "10");
    }
}

// FILE: tuple.rs
use crate::{input_generator::NextAttempt, InputGenerator, IntoInputGenerator};

fn cartesian_product_2<
    A: Iterator,
    B: Iterator,
    AMaker: Copy + Fn() -> A,
    BMaker: Copy + Fn() -> B,
>(
    a_maker: AMaker,
    b_maker: BMaker,
) -> impl Iterator<Item = (A::Item, B::Item)>
where
    A::Item: Clone,
    B::Item: Clone,
{
    a_maker().flat_map(move |a| b_maker().map(move |b| (a.clone(), b)))
}

fn cartesian_product_3<
    A: Iterator,
    B: Iterator,
    C: Iterator,
    AMaker: Copy + Fn() -> A,
    BMaker: Copy + Fn() -> B,
    CMaker: Copy + Fn() -> C,
>(
    a_maker: AMaker,
    b_maker: BMaker,
    c_maker: CMaker,
) -> impl Iterator<Item = (A::Item, B::Item, C::Item)>
where
    A::Item: Clone,
    B::Item: Clone,
    C::Item: Clone,
{
    cartesian_product_2(a_maker, move || cartesian_product_2(b_maker, c_maker))
        .map(|(a, (b, c))| (a, b, c))
}

#[non_exhaustive]
#[doc(hidden)]
pub struct TupleGen<A, B, C>(A, B, C);

impl<
        A,
        B,
        C,
        IntoGenA: IntoInputGenerator<A>,
        IntoGenB: IntoInputGenerator<B>,
        IntoGenC: IntoInputGenerator<C>,
    > IntoInputGenerator<(A, B, C)> for (IntoGenA, IntoGenB, IntoGenC)
{
    type IntoInputGenerator = TupleGen<
        IntoGenA::IntoInputGenerator,
        IntoGenB::IntoInputGenerator,
        IntoGenC::IntoInputGenerator,
    >;

    fn into_input_generator(self) -> Self::IntoInputGenerator {
        TupleGen(
            self.0.into_input_generator(),
            self.1.into_input_generator(),
            self.2.into_input_generator(),
        )
    }
}

#[allow(unused)] // TODO: WIP
pub struct TupleHist3<A, B, C> {
    inner_histories: (A, B, C),
    phase: TupleHist3Phase,
}

#[allow(unused)] // TODO: WIP
enum TupleHist3Phase {
    ElementWiseFirstPass { curr_index: u8 },
    AllTogetherFirstPass,
    AllPairs { first_index: u8, second_index: u8 },
    ElementWiseSecondPass { curr_index: u8 },
    AllTogetherSecondPass,
}

impl<
        A,
        B,
        C,
        GenA: InputGenerator<Input = A>,
        GenB: InputGenerator<Input = B>,
        GenC: InputGenerator<Input = C>,
    > InputGenerator for TupleGen<GenA, GenB, GenC>
{
    type Input = (A, B, C);

    type InputSource = (GenA::InputSource, GenB::InputSource, GenC::InputSource);

    type History = TupleHist3<GenA::History, GenB::History, GenC::History>;

    fn cardinality(&self) -> Option<usize> {
        Some(
            1usize
                .checked_mul(self.0.cardinality()?)?
                .checked_mul(self.1.cardinality()?)?
                .checked_mul(self.2.cardinality()?)?,
        )
    }

    fn exhaustive(&self) -> impl Iterator<Item = Self::InputSource> {
        cartesian_product_3(
            || self.0.exhaustive(),
            || self.1.exhaustive(),
            || self.2.exhaustive(),
        )
    }

    fn adversarial_count(&self) -> Option<usize> {
        Some(
            1usize
                .checked_mul(self.0.adversarial_count()?)?
                .checked_mul(self.1.adversarial_count()?)?
                .checked_mul(self.2.adversarial_count()?)?,
        )
    }

    fn adversarial(&self) -> impl Iterator<Item = Self::InputSource> {
        cartesian_product_3(
            || self.0.adversarial(),
            || self.1.adversarial(),
            || self.2.adversarial(),
        )
    }

    fn sample(&self, rng: &mut (impl rand::Rng + ?Sized)) -> Self::InputSource {
        (self.0.sample(rng), self.1.sample(rng), self.2.sample(rng))
    }

    fn new_history(&self, failing_input: Self::InputSource) -> Self::History {
        TupleHist3 {
            inner_histories: (
                self.0.new_history(failing_input.0),
                self.1.new_history(failing_input.1),
                self.2.new_history(failing_input.2),
            ),
            phase: TupleHist3Phase::ElementWiseFirstPass { curr_index: 0 },
        }
    }

    fn next_input(
        &self,
        rng: &mut impl rand::Rng,
        history: &Self::History,
    ) -> NextAttempt<Self::InputSource> {
        match history.phase {
            TupleHist3Phase::ElementWiseFirstPass { curr_index } => match curr_index {
                0 => match self.0.next_input(rng, &history.inner_histories.0) {
                    NextAttempt::Done => todo!(),
                    NextAttempt::InfoGathering(attempt) => NextAttempt::InfoGathering((
                        attempt,
                        self.1.current_simplest_failing(&history.inner_histories.1),
                        self.2.current_simplest_failing(&history.inner_histories.2),
                    )),
                    NextAttempt::ShrinkAttempt(attempt) => NextAttempt::ShrinkAttempt((
                        attempt,
                        self.1.current_simplest_failing(&history.inner_histories.1),
                        self.2.current_simplest_failing(&history.inner_histories.2),
                    )),
                },
                1 => todo!(),
                2 => todo!(),
                _ => unreachable!(),
            },
            TupleHist3Phase::AllTogetherFirstPass => todo!(),
            TupleHist3Phase::AllPairs {
                first_index,
                second_index,
            } => todo!(),
            TupleHist3Phase::ElementWiseSecondPass { curr_index } => todo!(),
            TupleHist3Phase::AllTogetherSecondPass => todo!(),
        }
    }

    fn current_simplest_failing(&self, history: &Self::History) -> Self::InputSource {
        (
            self.0.current_simplest_failing(&history.inner_histories.0),
            self.1.current_simplest_failing(&history.inner_histories.1),
            self.2.current_simplest_failing(&history.inner_histories.2),
        )
    }

    fn update_history(
        &self,
        _history: &mut Self::History,
        _shrinkable_input: Self::InputSource,
        _test_passed: bool,
    ) {
        todo!()
    }

    fn generate_observations(&self, _history: Self::History) -> Vec<crate::report::Observation> {
        todo!()
    }

    fn create_input(&self, input_source: Self::InputSource) -> Self::Input {
        (
            self.0.create_input(input_source.0),
            self.1.create_input(input_source.1),
            self.2.create_input(input_source.2),
        )
    }
}