/*
 * Look for TODO comments in this file for suggestions on how to implement
 * your parser.
 */
package parser;

import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

import javax.lang.model.util.ElementScanner6;

import lexer.ExprLexer;
import lexer.ParenLexer;
import lexer.SimpleLexer;
import lexer.TinyLexer;
import org.antlr.v4.runtime.*;

/**
 *
 */
public class Parser {

  final Grammar grammar;

  /**
   * All states in the parser.
   */
  private final States states;

  /**
   * Action table for bottom-up parsing. Accessed as
   * actionTable.get(state).get(terminal). You may replace
   * the Integer with a State class if you choose.
   */
  private final HashMap<Integer, HashMap<String, Action>> actionTable;
  /**
   * Goto table for bottom-up parsing. Accessed as
   * gotoTable.get(state).get(nonterminal).
   * You may replace the Integers with State classes if you choose.
   */
  private final HashMap<Integer, HashMap<String, Integer>> gotoTable;

  public Parser(String grammarFilename) throws IOException {
    actionTable = new HashMap<>();
    gotoTable = new HashMap<>();

    grammar = new Grammar(grammarFilename);

    states = new States();

    initializeStates();
    constructActionTable();
    constructGotoTable();
  }

  private int findState(State state) {
    for (State f_state : states.getStateList())
      if (f_state.equals(state))
        return f_state.getName();
    return -1;
  }

  private void constructActionTable() {
    for (State state : states.getStateList()) {
      int name = state.getName();
      actionTable.computeIfAbsent(name, k -> new HashMap<String, Action>());

      for (Item item : state.getItems()) {
        String nextSymbol = item.getNextSymbol();

        if (nextSymbol != null && grammar.isTerminal(nextSymbol))
          actionTable.get(name).put(nextSymbol, Action.createShift(findState(Parser.GOTO(state, nextSymbol, grammar))));
        else if (nextSymbol == null) {
          Rule rule = item.getRule();
          if (rule.equals(grammar.startRule) && item.getLookahead() == Util.EOF)
            actionTable.get(name).put(Util.EOF, Action.createAccept());
          else
            actionTable.get(name).put(item.getLookahead(), Action.createReduce(item.getRule()));
        }
      }
    }
  }

  private void constructGotoTable() {
    for (State state : states.getStateList()) {
      int name = state.getName();
      gotoTable.computeIfAbsent(name, k -> new HashMap<String, Integer>());
      for (String nonTerminal : grammar.nonterminals) {
        int goTo = findState(Parser.GOTO(state, nonTerminal, grammar));
        if (goTo != -1)
          gotoTable.get(name).put(nonTerminal, goTo);
      }
    }
  }

  private void initializeStates() {
    int stateId = 0;

    State closure = Parser.computeClosure(new Item(grammar.startRule, 0, Util.EOF), grammar);
    closure.setName(stateId++);
    states.addState(closure);

    boolean added = true;
    while (added) {
      added = false;

      ArrayList<State> stateList = new ArrayList<>(states.getStateList());
      for (State state : stateList)
        for (String symbol : grammar.symbols) {
          State goTo = Parser.GOTO(state, symbol, grammar);
          if (goTo.getItems().size() > 0 && states.addState(goTo)) {
            goTo.setName(stateId++);
            added = true;
          }
        }
    }
  }

  public States getStates() {
    return states;
  }

  static public State computeClosure(Item I, Grammar grammar) {
    State closure = new State(I);

    boolean added = true;
    while (added) {
      added = false;

      List<Item> items = new ArrayList<Item>(closure.getItems());

      for (Item item : items) {
        String symbol = item.getNextSymbol();

        if (grammar.isNonterminal(symbol)) {
          List<Rule> symbolRules = grammar.rules.stream().filter(rule -> rule.getLhs().equals(symbol))
              .collect(Collectors.toList());

          ArrayList<String> firstOperand = new ArrayList<>();
          if (item.getNextNextSymbol() != null) {
            List<String> rhs = item.getRule().getRhs();
            firstOperand.addAll(rhs.subList(item.getDot() + 1, rhs.size()));
          }
          firstOperand.add(item.getLookahead());

          HashSet<String> firstSet = new HashSet<>();
          for (String firstSymbol : firstOperand) {
            if (grammar.first.containsKey(firstSymbol)) {
              firstSet.addAll(grammar.first.get(firstSymbol));
              if (!grammar.first.get(firstSymbol).contains(Util.EPSILON))
                break;
            } else if (firstSymbol.equals(Util.EOF)) {
              firstSet.add(firstSymbol);
            }
          }

          for (Rule rule : symbolRules)
            for (String terminal : firstSet)
              added = closure.addItem(new Item(rule, 0, terminal)) || added;
        }
      }
    }

    return closure;
  }

  // This returns a new state that represents the transition from
  // the given state on the symbol X.
  static public State GOTO(State state, String X, Grammar grammar) {
    State gotoState = new State();

    for (Item item : state.getItems())
      if (item.getNextSymbol() != null && item.getNextSymbol().equals(X))
        gotoState.addItem(new Item(item.getRule(), item.getDot() + 1, item.getLookahead()));

    List<Item> gotoItems = new ArrayList<>(gotoState.getItems());
    for (Item gotoItem : gotoItems)
      for (Item closureItem : Parser.computeClosure(gotoItem, grammar).getItems())
        gotoState.addItem(closureItem);

    return gotoState;
  }

  // You will want to use StringBuilder. Another useful method will be
  // String.format: for
  // printing a value in the table, use
  // String.format("%8s", value)
  // How much whitespace you have shouldn't matter with regard to the tests, but
  // it will
  // help you debug if you can format it nicely.
  public String actionTableToString() {
    StringBuilder builder = new StringBuilder();

    List<String> terminalsPlusEof = new ArrayList<String>(grammar.terminals);
    terminalsPlusEof.add(Util.EOF);

    builder.append("state");
    for (String terminal : terminalsPlusEof)
      builder.append(String.format("%8s", terminal));
    builder.append("\n");

    for (State state : states.getStateList()) {
      int stateName = state.getName();
      HashMap<String, Action> actionRow = actionTable.get(stateName);
      builder.append(String.format("%5s", stateName));

      for (String terminal : terminalsPlusEof) {
        Action action = actionRow.get(terminal);
        builder.append(String.format("%8s", action != null ? action.toString() : ""));
      }
      builder.append("\n");
    }
    return builder.toString();
  }

  // You will want to use StringBuilder. Another useful method will be
  // String.format: for
  // printing a value in the table, use
  // String.format("%8s", value)
  // How much whitespace you have shouldn't matter with regard to the tests, but
  // it will
  // help you debug if you can format it nicely.
  public String gotoTableToString() {
    StringBuilder builder = new StringBuilder();

    builder.append("state");
    for (String nonTerminal : grammar.nonterminals)
      builder.append(String.format("%8s", nonTerminal));
    builder.append("\n");

    for (State state : states.getStateList()) {
      int stateName = state.getName();
      HashMap<String, Integer> gotoTableRow = gotoTable.get(stateName);
      builder.append(String.format("%5s", stateName));

      for (String nonTerminal : grammar.nonterminals) {
        if (gotoTableRow.containsKey(nonTerminal))
          builder.append(String.format("%8s", gotoTableRow.get(nonTerminal)));
        else
          builder.append(String.format("%8s", ""));
      }

      builder.append("\n");
    }
    return builder.toString();
  }

  // You should return a list of the actions taken.
  public List<Action> parse(Lexer scanner) throws ParserException {
    // tokens is the output from the scanner. It is the list of tokens
    // scanned from the input file.
    // To get the token type: v.getSymbolicName(t.getType())
    // To get the token lexeme: t.getText()
    ArrayList<? extends Token> tokens = new ArrayList<>(scanner.getAllTokens());
    Vocabulary v = scanner.getVocabulary();

    Stack<String> input = new Stack<>();
    Collections.reverse(tokens);
    input.add(Util.EOF);
    for (Token t : tokens) {
      input.push(v.getSymbolicName(t.getType()));
    }
    Collections.reverse(tokens);

    List<Action> actions = new ArrayList<>();

    Stack<State> stack = new Stack<>();
    stack.push(states.getState(0));

    String currentInput = input.pop();
    while (true) {
      State s = stack.peek();
      Action action = actionTable.get(s.getName()).get(currentInput);

      if (action == null)
        throw ParserException.create(tokens, input.size());

      actions.add(action);

      if (action.isShift()) {
        stack.push(states.getState(action.getState()));
        currentInput = input.pop();
      } else if (action.isReduce()) {
        Rule rule = action.getRule();
        for (int i = 0; i < rule.getRhs().size(); i++)
          stack.pop();

        s = stack.peek();
        stack.push(states.getState(gotoTable.get(s.getName()).get(rule.getLhs())));
      } else if (action.isAccept()) {
        break;
      } else {
        throw ParserException.create(tokens, input.size());
      }
    }

    return actions;
  }

  // -------------------------------------------------------------------
  // Convenience functions
  // -------------------------------------------------------------------

  public List<Action> parseFromFile(String filename) throws IOException, ParserException {
    // System.out.println("\nReading input file " + filename + "\n");
    final CharStream charStream = CharStreams.fromFileName(filename);
    Lexer scanner = scanFile(charStream);
    return parse(scanner);
  }

  public List<Action> parseFromString(String program) throws ParserException {
    Lexer scanner = scanFile(CharStreams.fromString(program));
    return parse(scanner);
  }

  private Lexer scanFile(CharStream charStream) {
    // We use ANTLR's scanner (lexer) to produce the tokens.
    Lexer scanner = null;
    switch (grammar.grammarName) {
      case "Simple":
        scanner = new SimpleLexer(charStream);
        break;
      case "Paren":
        scanner = new ParenLexer(charStream);
        break;
      case "Expr":
        scanner = new ExprLexer(charStream);
        break;
      case "Tiny":
        scanner = new TinyLexer(charStream);
        break;
      default:
        System.out.println("Unknown scanner");
        break;
    }

    return scanner;
  }

}