@file:Suppress("ForbiddenImport")

package viaduct.arbitrary.common

import io.kotest.common.ExperimentalKotest
import io.kotest.property.Arb
import io.kotest.property.Exhaustive
import io.kotest.property.Gen
import io.kotest.property.PropTestConfig
import io.kotest.property.PropertyContext
import io.kotest.property.PropertyTesting
import io.kotest.property.RandomSource
import io.kotest.property.Sample
import io.kotest.property.asSample
import io.kotest.property.checkAll
import java.lang.AssertionError
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.runBlocking
import viaduct.invariants.InvariantChecker

/** Convert an arb to an infinite [kotlin.sequences.Sequence] */
fun <T> Gen<T>.asSequence(rs: RandomSource): Sequence<T> = generate(rs).map { it.value }

/** return the smallest failing sample generated by this Gen within [iter] iterations */
fun <T> Gen<T>.minViolation(
    comparator: Comparator<T>,
    iter: Int? = null,
    property: (T) -> Boolean
): T? {
    val seq = when (this) {
        // for Exhaustives, calling asSequence(randomSource) will loop through the generator multiple times to fill
        // up to the requested iteration count. This is fine for Arb testing, but more repetitive for this simpler
        // method of finding the minimum violation.
        // As an alternative, just take all the Exhaustive values
        is Exhaustive -> this.values.asSequence()
        else -> asSequence(randomSource()).take(iter ?: PropertyTesting.defaultIterationCount)
    }
    return seq.filterNot(property).minWithOrNull(comparator)
}

/** assert that all values of an Arb pass the invariant checks defined by a provided function */
@OptIn(ExperimentalCoroutinesApi::class, ExperimentalKotest::class)
fun <T> Gen<T>.checkInvariants(
    iter: Int? = null,
    fn: PropertyContext.(T, InvariantChecker) -> Unit
) = runBlocking {
    val cfg = PropTestConfig(iterations = iter)
    checkAll(cfg) {
        val check = InvariantChecker()
        this.fn(it, check)
        check.assertEmpty("\n")
        markSuccess()
    }
}

/**
 * Flatten this Arb into an Arb of the inner item type.
 * The new Arb will return items in the same order as produced by the original Arb.
 */
fun <T> Arb<Iterable<T>>.flatten(): Arb<T> = Flatten(this)

private class Flatten<T>(val underlying: Arb<Iterable<T>>) : Arb<T>() {
    private var chunk: Iterator<T>? = null

    override fun edgecase(rs: RandomSource): T? = underlying.edgecase(rs)?.first()

    override fun sample(rs: RandomSource): Sample<T> {
        if (chunk == null || chunk?.hasNext() == false) {
            chunk = underlying.sample(rs).value.iterator()
        }
        return chunk!!.next().asSample()
    }
}

/**
 * Throw a property check failure.
 * If no seed value is provided, the message of the thrown exception will include the
 * current property testing seed value.
 */
fun failProperty(
    message: String,
    cause: Throwable? = null,
    seed: Long? = null
) {
    throw AssertionError(
        """
            |Property failed with seed ${seed ?: randomSource().seed}
            |$message
        """.trimMargin(),
        cause
    )
}
