#include <sstream>
#include <iomanip>
#include <cctype>
#include <algorithm>
#include <utility>
#include <set>
#include <array>
#include <limits.h>
#include "scanner.h"

Token::Token(Token::Kind kind, std::string lexeme):
  kind(kind), lexeme(std::move(lexeme)) {}

Token::Kind Token::getKind() const { return kind; }
const std::string &Token::getLexeme() const { return lexeme; }

std::ostream &operator<<(std::ostream &out, const Token &tok) {
  out << tok.getString() << '\n';
  return out;
}

const std::string Token::getString() const {
  std::string kind;
  switch (this->getKind()) {
    case Token::ID:         kind = "ID";        break;
    case Token::NUM:        kind = "NUM";       break;
    case Token::LPAREN:     kind = "LPAREN";    break;
    case Token::RPAREN:     kind = "RPAREN";    break;
    case Token::LBRACE:     kind = "LBRACE";    break;
    case Token::RBRACE:     kind = "RBRACE";    break;
    case Token::RETURN:     kind = "RETURN";    break;
    case Token::IF:         kind = "IF";        break;
    case Token::ELSE:       kind = "ELSE";      break;
    case Token::WHILE:      kind = "WHILE";     break;
    case Token::PRINTLN:    kind = "PRINTLN";   break;
    case Token::WAIN:       kind = "WAIN";      break;
    case Token::BECOMES:    kind = "BECOMES";   break;
    case Token::INT:        kind = "INT";       break;
    case Token::EQ:         kind = "EQ";        break;
    case Token::NE:         kind = "NE";        break;
    case Token::LT:         kind = "LT";        break;
    case Token::GT:         kind = "GT";        break;
    case Token::LE:         kind = "LE";        break;
    case Token::GE:         kind = "GE";        break;
    case Token::PLUS:       kind = "PLUS";      break;
    case Token::MINUS:      kind = "MINUS";     break;
    case Token::STAR:       kind = "STAR";      break;
    case Token::SLASH:      kind = "SLASH";     break;
    case Token::PCT:        kind = "PCT";       break;
    case Token::COMMA:      kind = "COMMA";     break;
    case Token::SEMI:       kind = "SEMI";      break;
    case Token::NEW:        kind = "NEW";       break;
    case Token::DELETE:     kind = "DELETE";    break;
    case Token::LBRACK:     kind = "LBRACK";    break;
    case Token::RBRACK:     kind = "RBRACK";    break;
    case Token::AMP:        kind = "AMP";       break;
    case Token::NUL:        kind = "NULL";      break;
    default: break;
  }

  return kind + ' ' + this->getLexeme();
}

int64_t Token::toNumber() const {
  std::istringstream iss;
  int64_t result;

  if (kind == NUM) {
    iss.str(lexeme);
  } else {
    // This should never happen if the user calls this function correctly
    return 0;
  }

  iss >> result;
  return result;
}

ScanningFailure::ScanningFailure(std::string message):
  message(std::move(message)) {}

const std::string &ScanningFailure::what() const { return message; }

/* Representation of a DFA, used to handle the scanning process.
 */
class AsmDFA {
  public:
    enum State {
      // States that are also kinds
      ID = 0,
      NUM,
      TOKEN,
      EQUAL,
      EXC,
      LG,
      ZERO,
      WHITESPACE,
      START,
      FAIL,

      // Hack to let this be used easily in arrays. This should always be the
      // final element in the enum, and should always point to the previous
      // element.

      LARGEST_STATE = FAIL
    };

  private:
    /* A set of all accepting states for the DFA.
     * Currently non-accepting states are not actually present anywhere
     * in memory, but a list can be found in the constructor.
     */
    std::set<State> acceptingStates;

    /*
     * The transition function for the DFA, stored as a map.
     */

    std::array<std::array<State, 128>, LARGEST_STATE + 1> transitionFunction;

    /*
     * Converts a state to a kind to allow construction of Tokens from States.
     * Throws an exception if conversion is not possible.
     */
    Token::Kind stateToKind(State s, const std::string &input) const {
      switch(s) {
        // Change to Respective Tokens { ID, NUM, TOKEN, EQUAL, LG, ZERO }
        case ID: {
          if (input == "wain") {
            return Token::WAIN;
          } else if (input == "int") {
            return Token::INT;
          } else if (input == "else") {
            return Token::ELSE;
          } else if (input == "while") {
            return Token::WHILE;
          } else if (input == "println") {
            return Token::PRINTLN;
          } else if (input == "return") {
            return Token::RETURN;
          } else if (input == "if") {
            return Token::IF;
          } else if (input == "NULL") {
            return Token::NUL;
          } else if (input == "new") {
            return Token::NEW;
          } else if (input == "delete") {
            return Token::DELETE;
          }
          return Token::ID;
        }
        case NUM: {
          std::istringstream iss;
          int64_t result;
          iss.str(input);
          iss >> result;
          if (result > INT_MAX) {
            throw ScanningFailure("ScanningFailure: Numeric literal out of range");
          }
          return Token::NUM;
        }
        case TOKEN: {
          if (input == "(") {
            return Token::LPAREN;
          } else if (input == ")") {
            return Token::RPAREN;
          } else if (input == "{") {
            return Token::LBRACE;
          } else if (input == "}") {
            return Token::RBRACE;
          } else if (input == "==") {
            return Token::EQ;
          } else if (input == "!=") {
            return Token::NE;
          } else if (input == "<=") {
            return Token::LE;
          } else if (input == ">=") {
            return Token::GE;
          } else if (input == "+") {
            return Token::PLUS;
          } else if (input == "-") {
            return Token::MINUS;
          } else if (input == "*") {
            return Token::STAR;
          } else if (input == "/") {
            return Token::SLASH;
          } else if (input == "%") {
            return Token::PCT;
          } else if (input == ",") {
            return Token::COMMA;
          } else if (input == ";") {
            return Token::SEMI;
          } else if (input == "[") {
            return Token::LBRACK;
          } else if (input == "]") {
            return Token::RBRACK;
          } else if (input == "&") {
            return Token::AMP;
          }
          throw ScanningFailure("ScanningFailure: Cannot convert state to kind. Lexeme: " + input);
        }
        case EQUAL: return Token::BECOMES;
        case LG: {
          if (input == ">") {
            return Token::GT;
          } else if (input == "<") {
            return Token::LT;
          }
          throw ScanningFailure("ScanningFailure: Cannot convert state to kind. Lexeme: " + input);
        }
        case ZERO:     return Token::NUM;
        case WHITESPACE: return Token::WHITESPACE;
        default: throw ScanningFailure("ScanningFailure: Cannot convert state to kind.");
      }
    }


  public:
    /* Tokenizes an input string according to the Simplified Maximal Munch
     * scanning algorithm.
     */
    std::vector<Token> simplifiedMaximalMunch(const std::string &input) const {
      std::vector<Token> result;

      State state = start();
      std::string munchedInput;

      for (std::string::const_iterator inputPosn = input.begin();
           inputPosn != input.end();) {

        State oldState = state;
        state = transition(state, *inputPosn);

        if (!failed(state)) {
          munchedInput += *inputPosn;
          oldState = state;

          ++inputPosn;
        }

        if (inputPosn == input.end() || failed(state)) {
          if (accept(oldState)) {
            result.push_back(Token(stateToKind(oldState, munchedInput), munchedInput));

            munchedInput = "";
            state = start();
          } else {
            if (failed(state)) {
              munchedInput += *inputPosn;
            }
            throw ScanningFailure("ScanningFailure: Simplified maximal munch failed on input: "
                                 + munchedInput);
          }
        }
      }

      return result;
    }

    /* Initializes the accepting states for the DFA.
     */
    AsmDFA() {
      acceptingStates = {ID, NUM, TOKEN, EQUAL, LG, ZERO, WHITESPACE};
      // Non-accepting states are EXC, START

      // Initialize transitions for the DFA
      for (size_t i = 0; i < transitionFunction.size(); ++i) {
        for (size_t j = 0; j < transitionFunction[0].size(); ++j) {
          transitionFunction[i][j] = FAIL;
        }
      }

      registerTransition(START, "(){}+-*/%,;[]&", TOKEN);
      registerTransition(START, "=", EQUAL);
      registerTransition(EQUAL, "=", TOKEN);
      registerTransition(START, "!", EXC);
      registerTransition(EXC, "=", TOKEN);
      registerTransition(START, "<>", LG);
      registerTransition(LG, "=", TOKEN);
      registerTransition(START, "0", ZERO);
      registerTransition(START, "123456789", NUM);
      registerTransition(NUM, isdigit, NUM);
      registerTransition(START, isalpha, ID);
      registerTransition(ID, isalnum, ID);
      registerTransition(START, isspace, WHITESPACE);
      registerTransition(WHITESPACE, isspace, WHITESPACE);
    }

    // Register a transition on all chars in chars
    void registerTransition(State oldState, const std::string &chars,
        State newState) {
      for (char c : chars) {
        transitionFunction[oldState][c] = newState;
      }
    }

    // Register a transition on all chars matching test
    // For some reason the cctype functions all use ints, hence the function
    // argument type.
    void registerTransition(State oldState, int (*test)(int), State newState) {

      for (int c = 0; c < 128; ++c) {
        if (test(c)) {
          transitionFunction[oldState][c] = newState;
        }
      }
    }

    /* Returns the state corresponding to following a transition
     * from the given starting state on the given character,
     * or a special fail state if the transition does not exist.
     */
    State transition(State state, char nextChar) const {
      return transitionFunction[state][nextChar];
    }

    /* Checks whether the state returned by transition
     * corresponds to failure to transition.
     */
    bool failed(State state) const { return state == FAIL; }

    /* Checks whether the state returned by transition
     * is an accepting state.
     */
    bool accept(State state) const {
      return acceptingStates.count(state) > 0;
    }

    /* Returns the starting state of the DFA
     */
    State start() const { return START; }
};

std::vector<Token> scan(const std::string &input) {
  static AsmDFA theDFA;

  std::vector<Token> tokens = theDFA.simplifiedMaximalMunch(input);

  // We need to:
  // * Throw exceptions for WORD tokens whose lexemes aren't ".word".
  // * Remove WHITESPACE and COMMENT tokens entirely.

  std::vector<Token> newTokens;

  for (auto &token : tokens) {
    if (token.getKind() != Token::WHITESPACE) {
      newTokens.push_back(token);
    }
  }

  return newTokens;
}