/*
 * Copyright (C) 2018-2024 University of Waterloo.
 *
 * This file is part of Perses.
 *
 * Perses is free software; you can redistribute it and/or modify it under the
 * terms of the GNU General Public License as published by the Free Software
 * Foundation; either version 3, or (at your option) any later version.
 *
 * Perses is distributed in the hope that it will be useful, but WITHOUT ANY
 * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along with
 * Perses; see the file LICENSE.  If not see <http://www.gnu.org/licenses/>.
 */
package org.perses.antlr.ast

import com.google.common.collect.ImmutableList
import org.perses.antlr.ast.RuleNameRegistry.RuleNameHandle
import org.perses.util.ast.Indent
import org.perses.util.exhaustive
import org.perses.util.toImmutableList
import org.perses.util.toImmutableMap
import java.io.PrintStream

class PersesGrammar(
  val grammarType: GrammarType,
  val grammarName: String,
  val channelDefinitions: PersesChannelDefinitionAst,
  val tokenSpecifications: PersesTokenSpecificationAst,
  val options: PersesGrammarOptionsAst,
  val namedActions: ImmutableList<PersesNamedAction>,
  val lexerRules: LexerRuleList,
  val parserRules: ImmutableList<PersesParserRuleAst>,
  val symbolTable: SymbolTable,
) : AbstractPersesAst() {

  enum class GrammarType {
    LEXER, PARSER, COMBINED
  }

  init {
    when (grammarType) {
      GrammarType.LEXER -> require(parserRules.isEmpty())
      GrammarType.PARSER -> require(lexerRules.flattenedLexerRules.isEmpty())
      else -> Unit
    }
  }

  val flattenedAllRules: ImmutableList<AbstractPersesRuleDefAst> = ImmutableList
    .builder<AbstractPersesRuleDefAst>()
    .addAll(lexerRules.flattenedLexerRules)
    .addAll(parserRules)
    .build()

  fun flattenedParserRulesSequence(): Sequence<Pair<RuleNameHandle, AbstractPersesRuleElement>> {
    return parserRules.asSequence()
      .flatMap { rule ->
        val body = rule.body
        return@flatMap if (body is PersesAlternativeBlockAst) {
          body.alternatives
            .asSequence()
            .map { alt -> rule.ruleNameHandle to alt }
        } else {
          sequenceOf(rule.ruleNameHandle to body)
        }
      }
  }

  val parserRuleNameList: ImmutableList<RuleNameHandle> by lazy {
    parserRules.asSequence()
      .map { it.ruleNameHandle }
      .toImmutableList()
  }

  private val allChildren = ImmutableList
    .builder<AbstractPersesAst>()
    .add(lexerRules)
    .addAll(parserRules)
    .build()

  private val ruleNameToRuleMap = flattenedAllRules
    .asSequence()
    .toImmutableMap(
      keyFunc = { it.ruleNameHandle },
      valueFunc = { it },
    )

  fun getRuleNameHandleOrThrow(ruleName: String): RuleNameHandle {
    return symbolTable.ruleNameRegistry.getOrThrow(ruleName)
  }

  fun copyWithNewName(newGrammarName: String): PersesGrammar {
    return PersesGrammar(
      grammarType,
      newGrammarName,
      channelDefinitions,
      tokenSpecifications,
      options,
      namedActions,
      lexerRules,
      parserRules,
      symbolTable,
    )
  }

  fun copyWithNewLexerRuleDefs(
    newLexerRules: LexerRuleList,
  ): PersesGrammar {
    return PersesGrammar(
      grammarType,
      grammarName,
      channelDefinitions,
      tokenSpecifications,
      options,
      namedActions,
      newLexerRules,
      parserRules,
      symbolTable,
    )
  }

  fun copyWithNewParserRuleDefs(
    newParserRules: ImmutableList<PersesParserRuleAst>,
  ): PersesGrammar {
    return PersesGrammar(
      grammarType,
      grammarName,
      channelDefinitions,
      tokenSpecifications,
      options,
      namedActions,
      lexerRules,
      newParserRules,
      symbolTable,
    )
  }

  // TODO: to be unit-tested.
  fun copyWithNewRuleDef(newDef: PersesParserRuleAst): PersesGrammar {
    val newParserRules = parserRules
      .asSequence()
      .map { existing: PersesParserRuleAst ->
        if (existing.ruleNameHandle == newDef.ruleNameHandle) {
          return@map newDef
        } else {
          return@map existing
        }
      }.toImmutableList()
    return PersesGrammar(
      grammarType,
      grammarName,
      channelDefinitions,
      tokenSpecifications,
      options,
      namedActions,
      lexerRules,
      newParserRules,
      symbolTable,
    )
  }

  fun getRuleDefinition(ruleName: String): AbstractPersesRuleDefAst? {
    val ruleNameHandle = symbolTable.ruleNameRegistry[ruleName] ?: return null
    return getRuleDefinition(ruleNameHandle)
  }

  fun getRuleDefinition(ruleNameHandle: RuleNameHandle): AbstractPersesRuleDefAst? {
    return ruleNameToRuleMap[ruleNameHandle]
  }

  override fun toSourceCode(stream: PrintStream, indent: Indent, multiLineMode: Boolean) {
    stream.println("// Auto-generated by Perses AST.")
    stream.println("// DO NOT MODIFY.")
    when (grammarType) {
      GrammarType.LEXER -> stream.print("lexer ")
      GrammarType.PARSER -> stream.print("parser ")
      GrammarType.COMBINED -> Unit
    }.exhaustive
    stream.printf("grammar %s;\n", grammarName)
    stream.print('\n')
    val emptyIndent = Indent.FOUR_CHAR_EMPTY_INDENT
    channelDefinitions.toSourceCode(stream, emptyIndent, SINGLE_LINE_MODE)
    tokenSpecifications.toSourceCode(stream, emptyIndent, SINGLE_LINE_MODE)
    options.toSourceCode(stream, emptyIndent, SINGLE_LINE_MODE)
    for (action in namedActions) {
      action.toSourceCode(stream, emptyIndent, SINGLE_LINE_MODE)
    }
    lexerRules.toSourceCode(stream, emptyIndent, MULTI_LINE_MODE)
    stream.println("\n")
    for (rule in parserRules) {
      rule.toSourceCode(stream, emptyIndent, MULTI_LINE_MODE)
      stream.print('\n')
    }
  }

  override val childCount: Int
    get() = allChildren.size

  override fun getChild(index: Int): AbstractPersesAst {
    return allChildren[index]
  }

  override val extraLabelForTreeStructurePrinting: String
    get() = super.extraLabelForTreeStructurePrinting + " grammar_name=" + grammarName

  override val tag = AstTag.GRAMMAR

  fun isEquivalent(other: PersesGrammar): Boolean {
    if (ruleNameToRuleMap.keys != other.ruleNameToRuleMap.keys) {
      return false
    }
    for (ruleNameHandle in ruleNameToRuleMap.keys) {
      if (!ruleNameToRuleMap[ruleNameHandle]!!
          .body
          .isEquivalent(other.ruleNameToRuleMap[ruleNameHandle]!!.body)
      ) {
        return false
      }
    }
    return true
  }

  fun computeAntlrBaseFileName() = "$grammarName.g4"

  companion object {
    private const val SINGLE_LINE_MODE = false
    private const val MULTI_LINE_MODE = true
  }
}
