package io.kotest.property.arbitrary

import io.kotest.property.Arb
import io.kotest.property.Gen

/**
 * Returns a [Arb] that generates randomly-chosen Chars. Custom characters can be generated by
 * providing CharRanges. Distribution will be even across the ranges of Chars.
 * For example:
 * Gen.char('A'..'C', 'D'..'E')
 * Ths will choose A, B, C, D, and E each 20% of the time.
 */
fun Arb.Companion.char(range: CharRange, vararg ranges: CharRange): Arb<Char> {
   return Arb.char(listOf(range) + ranges)
}

/**
 * Returns a [Arb] that generates randomly-chosen Chars. Custom characters can be generated by
 * providing a list of CharRanges. Distribution will be even across the ranges of Chars.
 * For example:
 * Gen.char(listOf('A'..'C', 'D'..'E')
 * Ths will choose A, B, C, D, and E each 20% of the time.
 *
 * If no parameter is given, ASCII characters will be generated.
 */
fun Arb.Companion.char(ranges: List<CharRange> = CharSets.BASIC_LATIN): Arb<Char> {

   require(ranges.all { !it.isEmpty() }) { "Ranges cannot be empty" }
   require(ranges.isNotEmpty()) { "List of ranges must have at least one range" }

   fun makeRangeWeightedGen(): Arb<CharRange> {
      val weightPairs = ranges.map { range ->
         val weight = range.last.code - range.first.code + 1
         Pair(weight, range)
      }
      return Arb.choose(weightPairs[0], weightPairs[1], *weightPairs.drop(2).toTypedArray())
   }

   // Convert the list of CharRanges into a weighted Gen in which
   // the ranges are chosen from the list using the length of the
   // range as the weight.
   val arbRange: Arb<CharRange> =
      if (ranges.size == 1) Arb.constant(ranges.first())
      else makeRangeWeightedGen()

   return arbRange.flatMap { charRange ->
      arbitrary { charRange.random(it.random) }
   }
}

/**
 * Returns an [Arb] that produces [CharArray]s where [length] produces the length of the arrays and
 * [content] produces the content of the arrays.
 */
fun Arb.Companion.charArray(length: Gen<Int>, content: Arb<Char>): Arb<CharArray> =
   toPrimitiveArray(length, content, Collection<Char>::toCharArray)

private object CharSets {
   val CONTROL = listOf('\u0000'..'\u001F', '\u007F'..'\u007F')
   val WHITESPACE = listOf('\u0020'..'\u0020', '\u0009'..'\u0009', '\u000A'..'\u000A')
   val BASIC_LATIN = listOf('\u0021'..'\u007E')
}
