///|
enum MarkovToken {
  Word(String)
  Punctuation(String)
}

///|
priv enum GenerationState {
  NeedStart
  Following(MarkovToken, MarkovToken)
}

///|
priv struct GenerationProgress {
  word_count : Int
  state : GenerationState
}

///|
pub struct SecondOrderMarkov {
  starts : Array[(MarkovToken, MarkovToken)]
  transitions : Map[String, Array[MarkovToken]]
  words : Array[String]
  default_punctuation : String
}

///|
fn MarkovToken::key(self : MarkovToken) -> String {
  match self {
    Word(word) => "w\{word.length()}:\{word}"
    Punctuation(punctuation) => "p\{punctuation.length()}:\{punctuation}"
  }
}

///|
fn markov_key(first : MarkovToken, second : MarkovToken) -> String {
  first.key() + "|" + second.key()
}

///|
fn push_transition(
  transitions : Map[String, Array[MarkovToken]],
  first : MarkovToken,
  second : MarkovToken,
  next : MarkovToken,
) -> Unit {
  let key = markov_key(first, second)
  let bucket = transitions.get_or_init(key, () => [])
  bucket.push(next)
}

///|
fn[T] choose(rand : @random.Rand, values : Array[T]) -> T? {
  let length = values.length()
  guard length != 0 else { None }
  Some(values[rand.int(limit=length)])
}

///|
fn append_word_token(
  tokens : Array[MarkovToken],
  words : Array[String],
  word : Word,
) -> Unit {
  let owned = word.to_owned()
  guard !owned.is_empty() else { () }
  tokens.push(Word(owned))
  words.push(owned)
}

///|
fn Sentence::tokens(
  self : Sentence,
  delimiter : Delimiter,
  words : Array[String],
) -> Array[MarkovToken] {
  let tokens = []
  for index, clause in self.0 {
    if index > 0 && delimiter.comma.length() > 0 {
      tokens.push(Punctuation(delimiter.comma))
    }
    clause.0.each(w => append_word_token(tokens, words, w))
  }
  let punctuation = delimiter.display_termination(self.1)
  if punctuation.length() > 0 {
    tokens.push(Punctuation(punctuation))
  }
  tokens
}

///|
fn tokens_have_word(tokens : Array[MarkovToken]) -> Bool {
  tokens.any(token => token is Word(_))
}

///|
pub fn SecondOrderMarkov::from_paragraph(
  delimiter : Delimiter,
  paragraph : Paragraph,
) -> SecondOrderMarkov {
  let starts = []
  let transitions : Map[String, Array[MarkovToken]] = Map([])
  let words = []
  let stream = []
  for sentence in paragraph.0 {
    let sentence_tokens = sentence.tokens(delimiter, words)
    if sentence_tokens.length() >= 2 && tokens_have_word(sentence_tokens) {
      starts.push((sentence_tokens[0], sentence_tokens[1]))
    }
    sentence_tokens.each(token => stream.push(token))
  }
  for i in 0..<(stream.length() - 2) {
    push_transition(transitions, stream[i], stream[i + 1], stream[i + 2])
  }
  { starts, transitions, words, default_punctuation: delimiter.period }
}

///|
pub fn SecondOrderMarkov::from_text(
  delimiter : Delimiter,
  corpus : StringView,
) -> SecondOrderMarkov {
  let (paragraph, _) = delimiter.paragraph(corpus)
  SecondOrderMarkov::from_paragraph(delimiter, paragraph)
}

///|
fn render_tokens(delimiter : Delimiter, tokens : Array[MarkovToken]) -> String {
  let mut text = ""
  for token in tokens {
    match token {
      Word(word) if text.length() == 0 => text = word
      Word(word) => text = text + delimiter.space + word
      Punctuation(punctuation) => text = text + punctuation
    }
  }
  text
}

///|
fn append_token(
  generated : Array[MarkovToken],
  token : MarkovToken,
  count : Int,
  limit : Int,
) -> Int? {
  match token {
    Word(_) if count >= limit => None
    Word(_) => {
      generated.push(token)
      Some(count + 1)
    }
    Punctuation(_) => {
      generated.push(token)
      Some(count)
    }
  }
}

///|
fn append_start(
  generated : Array[MarkovToken],
  first : MarkovToken,
  second : MarkovToken,
  count : Int,
  limit : Int,
) -> GenerationProgress {
  guard append_token(generated, first, count, limit) is Some(first_count) else {
    { word_count: count, state: NeedStart }
  }
  guard append_token(generated, second, first_count, limit)
    is Some(second_count) else {
    { word_count: first_count, state: NeedStart }
  }
  { word_count: second_count, state: Following(first, second) }
}

///|
fn SecondOrderMarkov::append_fallback_word(
  self : SecondOrderMarkov,
  generated : Array[MarkovToken],
  count : Int,
  limit : Int,
  rand : @random.Rand,
) -> GenerationProgress {
  guard choose(rand, self.words) is Some(word) else {
    { word_count: limit, state: NeedStart }
  }
  generated.push(Word(word))
  { word_count: count + 1, state: NeedStart }
}

///|
fn SecondOrderMarkov::begin_generation(
  self : SecondOrderMarkov,
  generated : Array[MarkovToken],
  count : Int,
  limit : Int,
  rand : @random.Rand,
) -> GenerationProgress {
  guard choose(rand, self.starts) is Some((first, second)) else {
    self.append_fallback_word(generated, count, limit, rand)
  }
  append_start(generated, first, second, count, limit)
}

///|
fn SecondOrderMarkov::follow_transition(
  self : SecondOrderMarkov,
  generated : Array[MarkovToken],
  previous : MarkovToken,
  current : MarkovToken,
  count : Int,
  limit : Int,
  rand : @random.Rand,
) -> GenerationProgress {
  guard self.transitions.get(markov_key(previous, current)) is Some(next_tokens) else {
    { word_count: count, state: NeedStart }
  }
  guard choose(rand, next_tokens) is Some(next) else {
    { word_count: count, state: NeedStart }
  }
  guard append_token(generated, next, count, limit) is Some(next_count) else {
    { word_count: count, state: NeedStart }
  }
  { word_count: next_count, state: Following(current, next) }
}

///|
fn SecondOrderMarkov::can_generate(
  self : SecondOrderMarkov,
  word_limit : Int,
) -> Bool {
  word_limit > 0 && !self.words.is_empty()
}

///|
pub fn SecondOrderMarkov::generate(
  self : SecondOrderMarkov,
  delimiter : Delimiter,
  words~ : Int,
  rand? : @random.Rand = @random.Rand::chacha8(seed=generate_time_based_seed()),
) -> String {
  guard self.can_generate(words) else { "" }
  let generated = []
  let mut progress = { word_count: 0, state: NeedStart }
  while progress.word_count < words {
    progress = match progress.state {
      NeedStart =>
        self.begin_generation(generated, progress.word_count, words, rand)
      Following(previous, current) =>
        self.follow_transition(
          generated,
          previous,
          current,
          progress.word_count,
          words,
          rand,
        )
    }
  }
  match generated.last() {
    Some(Punctuation(_)) => ()
    Some(_) => generated.push(Punctuation(self.default_punctuation))
    None => ()
  }
  render_tokens(delimiter, generated)
}