package cs241e.scanparse

import DFAs.Token

object Grammars {

  /** A representation of a production rule in a context-free grammar. */
  case class Production(lhs: String, rhs: Seq[String]) {
    override def toString: String = lhs + " " + rhs.mkString(" ")
  }

  /** A representation of a context-free grammar. */
  case class Grammar(productions: Seq[Production]) {
    val nonTerminals = productions.map(_.lhs).toSet
    val symbols = nonTerminals ++ productions.flatMap(_.rhs).toSet
    val terminals = symbols -- nonTerminals
    val start = productions.head.lhs

    /** Given a non-terminal symbol, returns the set of production rules that have the symbol as their left-hand side.
      **/
    val productionsExpanding: Map[String, Seq[Production]] =
      productions.groupBy(_.lhs).withDefaultValue(Seq())

    override def toString: String = productions.mkString("\n")
  }

  /** A representation of a node of a parse tree.
    *
    * The grammar symbol corresponding to the tree node is in `lhs.kind`. If the tree node is a leaf node representing
    * a terminal (token), the lexeme of the token is in `lhs.lexeme`; otherwise, lhs.lexeme is the empty string.
    */
  class Tree(val lhs: Token, val children: Seq[Tree] = Seq.empty) {
    def production: String = lhs.kind + " " + children.map(_.lhs.kind).mkString(" ")
    def show(indent: Int = 0): String = " " * indent + lhs + "\n" + children.map(_.show(indent + 1)).mkString
    override def toString: String = show(0)
  }

  /** Generates a `Grammar` from a textual description. Each non-empty line of the input is considered a production
    * rule of the grammar. The symbols of the production rule are space-separated words on the line. The first word
    * is the lhs of the production rule, and the remaining words are the rhs. The lhs of the first production rule
    * is considered the start symbol of the grammar.
    */
  def parseGrammar(input: String): Grammar = {
    def parseLine(line: String): Option[Production] = {
      val tokens = line.trim.split(" ").filterNot(_.isEmpty)
      if (tokens.isEmpty) None else Some(Production(tokens.head, tokens.tail))
    }
    val productions = input.split("\n").flatMap(parseLine)
    Grammar(productions)
  }
}