#include <tree_sitter/parser.h>
#include <string>
#include <cwctype>

namespace {

using std::string;

enum TokenType {
  DOLLAR_QUOTED_STRING_TAG,
  DOLLAR_QUOTED_STRING_CONTENT,
  DOLLAR_QUOTED_STRING_END_TAG,
};

struct Scanner {
  string dollar_quoted_string_tag;
  string current_leading_word;
  bool dollar_quoted_string_started;

  void skip(TSLexer *lexer) {
    lexer->advance(lexer, true);
  }

  void advance(TSLexer *lexer) {
    lexer->advance(lexer, false);
  }

  unsigned serialize(char *buffer) {
    if (dollar_quoted_string_tag.length() + 1>= TREE_SITTER_SERIALIZATION_BUFFER_SIZE) return 0;
    buffer[0] = dollar_quoted_string_started;
    dollar_quoted_string_tag.copy(&buffer[1], dollar_quoted_string_tag.length());
    return dollar_quoted_string_tag.length() + 1;
  }

  void deserialize(const char *buffer, unsigned length) {
    if (length == 0) {
      dollar_quoted_string_started = false;
      dollar_quoted_string_tag.clear();
    } else {
      dollar_quoted_string_started = buffer[0];
      dollar_quoted_string_tag.assign(&buffer[1], &buffer[length]);
    }
  }

  bool scan_dollar_quoted_string_content(TSLexer *lexer) {
    unsigned long int pos = 0;
    lexer->result_symbol = DOLLAR_QUOTED_STRING_CONTENT;
    lexer->mark_end(lexer);

    for (;;) {
        if (lexer->lookahead == '\0') {
          return false;
        } else if (lexer->lookahead == dollar_quoted_string_tag[pos]) {
            if (pos == dollar_quoted_string_tag.length() - 1) {
              return true;
            } else if (pos == 0) {
              lexer->result_symbol = DOLLAR_QUOTED_STRING_CONTENT;
              lexer->mark_end(lexer);
            }
            pos++;
            advance(lexer);
        } else if (pos != 0) {
            pos = 0;
        } else {
            advance(lexer);
        }
    }
  }

  bool scan_dollar_quoted_string_tag(TSLexer *lexer) {
    while (iswspace(lexer->lookahead)) skip(lexer);

    dollar_quoted_string_tag.clear();
    if (lexer->lookahead == '$') {
      dollar_quoted_string_tag += lexer->lookahead;
      advance(lexer);
    } else {
      return false;
    }

    while (iswalpha(lexer->lookahead)) {
      dollar_quoted_string_tag += lexer->lookahead;
      advance(lexer);
    }

    if (lexer->lookahead == '$') {
      dollar_quoted_string_tag += lexer->lookahead;
      advance(lexer);
      dollar_quoted_string_started = true;
      return true;
    }

    return false;
  }

  bool scan_dollar_quoted_string_end_tag(TSLexer *lexer) {
    current_leading_word.clear();
    while (
      lexer->lookahead != '\0' &&
      current_leading_word.length() < dollar_quoted_string_tag.length()
    ) {
      current_leading_word += lexer->lookahead;
      advance(lexer);
    }
    return current_leading_word == dollar_quoted_string_tag;
  }

  bool scan(TSLexer *lexer, const bool *valid_symbols) {
    if (valid_symbols[DOLLAR_QUOTED_STRING_TAG] && !dollar_quoted_string_started) {
      return scan_dollar_quoted_string_tag(lexer);
    }

    if (valid_symbols[DOLLAR_QUOTED_STRING_CONTENT] && dollar_quoted_string_started) {
      return scan_dollar_quoted_string_content(lexer);
    }

    if (valid_symbols[DOLLAR_QUOTED_STRING_END_TAG] && dollar_quoted_string_started) {
      if (scan_dollar_quoted_string_end_tag(lexer)) {
        dollar_quoted_string_started = false;
        lexer->result_symbol = DOLLAR_QUOTED_STRING_END_TAG;
        return true;
      };
    }

    return false;
  }

};

}

extern "C" {

void *tree_sitter_sql_external_scanner_create() {
  return new Scanner();
}

bool tree_sitter_sql_external_scanner_scan(void *payload, TSLexer *lexer,
                                            const bool *valid_symbols) {
  Scanner *scanner = static_cast<Scanner *>(payload);
  return scanner->scan(lexer, valid_symbols);
}

unsigned tree_sitter_sql_external_scanner_serialize(void *payload, char *state) {
  Scanner *scanner = static_cast<Scanner *>(payload);
  return scanner->serialize(state);
}

void tree_sitter_sql_external_scanner_deserialize(void *payload, const char *state, unsigned length) {
  Scanner *scanner = static_cast<Scanner *>(payload);
  scanner->deserialize(state, length);
}

void tree_sitter_sql_external_scanner_destroy(void *payload) {
  Scanner *scanner = static_cast<Scanner *>(payload);
  delete scanner;
}

}