#include "tokenizer.hpp"

#include "duckdb/common/printer.hpp"
#include "duckdb/common/string_util.hpp"

namespace duckdb {

BaseTokenizer::BaseTokenizer(const string &sql, vector<MatcherToken> &tokens) : sql(sql), tokens(tokens) {
}

static bool OperatorEquals(const char *str, const char *op, idx_t len, idx_t &op_len) {
	for (idx_t i = 0; i < len; i++) {
		if (str[i] != op[i]) {
			return false;
		}
	}
	op_len = len;
	return true;
}

bool BaseTokenizer::IsSpecialOperator(idx_t pos, idx_t &op_len) const {
	const char *op_start = sql.c_str() + pos;
	if (pos + 2 < sql.size()) {
		if (OperatorEquals(op_start, "->>", 3, op_len)) {
			return true;
		}
	}
	if (pos + 1 >= sql.size()) {
		// 2-byte operators are out-of-bounds
		return false;
	}
	if (OperatorEquals(op_start, "::", 2, op_len)) {
		return true;
	}
	if (OperatorEquals(op_start, ":=", 2, op_len)) {
		return true;
	}
	if (OperatorEquals(op_start, "->", 2, op_len)) {
		return true;
	}
	if (OperatorEquals(op_start, "**", 2, op_len)) {
		return true;
	}
	if (OperatorEquals(op_start, "//", 2, op_len)) {
		return true;
	}
	return false;
}

bool BaseTokenizer::IsSingleByteOperator(char c) {
	switch (c) {
	case '(':
	case ')':
	case '{':
	case '}':
	case '[':
	case ']':
	case ',':
	case '?':
	case '$':
	case '-':
	case '#':
		return true;
	default:
		return false;
	}
}

bool BaseTokenizer::CharacterIsInitialNumber(char c) {
	if (c >= '0' && c <= '9') {
		return true;
	}
	return c == '.';
}

bool BaseTokenizer::CharacterIsSpecialStringCharacter(char c) {
	if (c == 'N' || c == 'n') {
		return true;
	}
	if (c == 'X' || c == 'x') {
		return true;
	}
	if (c == 'E' || c == 'e') {
		return true;
	}
	if (c == 'B' || c == 'b') {
		return true;
	}
	return false;
}

bool BaseTokenizer::CharacterIsNumber(char c) {
	if (CharacterIsInitialNumber(c)) {
		return true;
	}
	switch (c) {
	case 'e': // exponents
	case 'E':
	case '_':
		return true;
	default:
		return false;
	}
}

bool BaseTokenizer::CharacterIsScientific(char c) {
	switch (c) {
	case 'e':
	case 'E':
		return true;
	default:
		return false;
	}
}

bool BaseTokenizer::CharacterIsControlFlow(char c) {
	switch (c) {
	case '\'':
	case '-':
	case ';':
	case '"':
	case '.':
		return true;
	default:
		return false;
	}
}

bool BaseTokenizer::CharacterIsKeyword(char c) {
	if (IsSingleByteOperator(c)) {
		return false;
	}
	if (StringUtil::CharacterIsOperator(c)) {
		return false;
	}
	if (StringUtil::CharacterIsSpace(c)) {
		return false;
	}
	if (CharacterIsControlFlow(c)) {
		return false;
	}
	return true;
}

bool BaseTokenizer::CharacterIsOperator(char c) {
	if (IsSingleByteOperator(c)) {
		return false;
	}
	if (CharacterIsControlFlow(c)) {
		return false;
	}
	return StringUtil::CharacterIsOperator(c);
}

void BaseTokenizer::PushToken(idx_t start, idx_t end) {
	if (start >= end) {
		return;
	}
	string last_token = sql.substr(start, end - start);
	tokens.emplace_back(std::move(last_token), start);
}

// Valid characters can be between A-Z, a-z, 0-9, underscore, or \200 - \377
// Note: 0-9 are only valid after the first character. Callers are expected to validate that before calling this
// function.
bool BaseTokenizer::IsValidDollarTagCharacter(char c) {
	if (c >= 'A' && c <= 'Z') {
		return true;
	}
	if (c >= 'a' && c <= 'z') {
		return true;
	}
	if (c >= '0' && c <= '9') {
		return true;
	}
	if (c == '_') {
		return true;
	}
	if (c >= '\200' && c <= '\377') {
		return true;
	}
	return false;
}

bool BaseTokenizer::TokenizeInput() {
	auto state = TokenizeState::STANDARD;
	idx_t last_pos = 0;
	string dollar_quote_marker;
	for (idx_t i = 0; i < sql.size(); i++) {
		auto c = sql[i];
		switch (state) {
		case TokenizeState::STANDARD:
			if (c == '\'') {
				state = TokenizeState::STRING_LITERAL;
				last_pos = i;
				break;
			}
			if (c == '"') {
				state = TokenizeState::QUOTED_IDENTIFIER;
				last_pos = i;
				break;
			}
			if (c == ';') {
				// end of statement
				OnStatementEnd(i);
				last_pos = i + 1;
				break;
			}
			if (c == '$') {
				// Dollar-quoted string statement
				if (i + 1 >= sql.size()) {
					// We need more than a single dollar
					break;
				}
				if (sql[i + 1] >= '0' && sql[i + 1] <= '9') {
					// $[numeric] is a parameter, not a dollar-quoted string
					tokens.emplace_back(string(1, c), i);
					break;
				}
				// Dollar-quoted string or collabel parameter ($collabel)
				last_pos = i;
				// Scan until next $
				idx_t next_dollar = 0;
				for (idx_t idx = i + 1; idx < sql.size(); idx++) {
					if (sql[idx] == '$') {
						next_dollar = idx;
						break;
					}
					if (!IsValidDollarTagCharacter(sql[idx])) {
						break;
					}
				}
				if (next_dollar == 0) {
					// Collabel parameter ($collabel)
					tokens.emplace_back(string(1, c), i);
					break;
				}
				state = TokenizeState::DOLLAR_QUOTED_STRING;
				last_pos = i;
				i = next_dollar;
				if (i < sql.size()) {
					// Found a complete marker, store it.
					idx_t marker_start = last_pos + 1;
					dollar_quote_marker = string(sql.begin() + marker_start, sql.begin() + i);
				}
				break;
			}
			if (c == '-' && i + 1 < sql.size() && sql[i + 1] == '-') {
				i++;
				state = TokenizeState::SINGLE_LINE_COMMENT;
				break;
			}
			if (c == '/' && i + 1 < sql.size() && sql[i + 1] == '*') {
				i++;
				state = TokenizeState::MULTI_LINE_COMMENT;
				break;
			}
			if (StringUtil::CharacterIsSpace(c)) {
				// space character - skip
				last_pos = i + 1;
				break;
			}
			idx_t op_len;
			if (IsSpecialOperator(i, op_len)) {
				// special operator - push the special operator
				tokens.emplace_back(sql.substr(i, op_len), last_pos);
				i += op_len - 1;
				last_pos = i + 1;
				break;
			}
			if (IsSingleByteOperator(c)) {
				// single-byte operator - directly push the token
				tokens.emplace_back(string(1, c), last_pos);
				last_pos = i + 1;
				break;
			}
			if (CharacterIsInitialNumber(c)) {
				// parse a numeric literal
				state = TokenizeState::NUMERIC;
				last_pos = i;
				break;
			}
			if (CharacterIsSpecialStringCharacter(c)) {
				// Look ahead to see if a quote follows
				if (i + 1 < sql.size() && sql[i + 1] == '\'') {
					state = TokenizeState::STRING_LITERAL;
					last_pos = i;
					i++;
					break;
				}
			}
			if (StringUtil::CharacterIsOperator(c)) {
				state = TokenizeState::OPERATOR;
				last_pos = i;
				break;
			}
			state = TokenizeState::KEYWORD;
			last_pos = i;
			break;
		case TokenizeState::NUMERIC:
			// Check for "always allowed" numeric characters
			if (CharacterIsInitialNumber(c)) {
				break; // Continue tokenizing
			}
			// Allow underscore only when immediately followed by a digit (no consecutive underscores)
			if (c == '_' && i + 1 < sql.size() && CharacterIsInitialNumber(sql[i + 1])) {
				break;
			}

			// Check for scientific notation marker
			if (CharacterIsScientific(c)) {
				// (e.g., "1ee5" is invalid)
				if (!CharacterIsScientific(sql[i - 1])) {
					break; // Valid 'e' or 'E', continue
				}
			}

			// Check for '+' or '-'
			if (c == '+' || c == '-') {
				if (CharacterIsScientific(sql[i - 1])) {
					break; // Valid, e.g., "1e-5". Continue.
				}
				// Invalid, e.g., "1+5" or "1e5+5". Fall through to stop.
			}

			// --- End of number ---
			// The character 'c' is not a valid part of the number.
			// Stop tokenizing and backtrack as per your original logic.
			while (!CharacterIsInitialNumber(sql[i - 1])) {
				i--;
			}
			PushToken(last_pos, i);
			state = TokenizeState::STANDARD;
			last_pos = i;
			i--;
			break;
		case TokenizeState::OPERATOR:
			// operator literal - check if this is still an operator
			if (!CharacterIsOperator(c)) {
				// Apply PostgreSQL trimming rule: an operator cannot end in '+' unless
				// it contains at least one of: ~ ! @ # % ^ & | ` ?
				idx_t end_pos = i;
				bool has_special = false;
				for (idx_t j = last_pos; j < end_pos; j++) {
					char oc = sql[j];
					if (oc == '~' || oc == '!' || oc == '@' || oc == '#' || oc == '%' || oc == '^' || oc == '&' ||
					    oc == '|' || oc == '`' || oc == '?') {
						has_special = true;
						break;
					}
				}
				if (!has_special) {
					while (end_pos > last_pos && sql[end_pos - 1] == '+') {
						end_pos--;
					}
				}
				PushToken(last_pos, end_pos);
				// Push any trimmed '+' characters as individual tokens
				for (idx_t j = end_pos; j < i; j++) {
					tokens.emplace_back(string(1, sql[j]), j);
				}
				state = TokenizeState::STANDARD;
				last_pos = i;
				i--;
			}
			break;
		case TokenizeState::KEYWORD:
			// keyword - check if this is still a keyword
			if (!CharacterIsKeyword(c)) {
				// not a keyword - return to standard state
				PushToken(last_pos, i);
				state = TokenizeState::STANDARD;
				last_pos = i;
				i--;
			}
			break;
		case TokenizeState::STRING_LITERAL:
			if (c == '\'') {
				if (i + 1 < sql.size() && sql[i + 1] == '\'') {
					// escaped - skip escape
					i++;
				} else {
					PushToken(last_pos, i + 1);
					last_pos = i + 1;
					state = TokenizeState::STANDARD;
				}
			}
			break;
		case TokenizeState::QUOTED_IDENTIFIER:
			if (c == '"') {
				if (i + 1 < sql.size() && sql[i + 1] == '"') {
					// escaped - skip escape
					i++;
				} else {
					PushToken(last_pos, i + 1);
					last_pos = i + 1;
					state = TokenizeState::STANDARD;
				}
			}
			break;
		case TokenizeState::SINGLE_LINE_COMMENT:
			if (c == '\n' || c == '\r') {
				last_pos = i + 1;
				state = TokenizeState::STANDARD;
			}
			break;
		case TokenizeState::MULTI_LINE_COMMENT:
			if (c == '*' && i + 1 < sql.size() && sql[i + 1] == '/') {
				i++;
				last_pos = i + 1;
				state = TokenizeState::STANDARD;
			}
			break;
		case TokenizeState::DOLLAR_QUOTED_STRING: {
			// Dollar-quoted string -- all that will get us out is a $[marker]$
			if (c != '$') {
				break;
			}
			if (i + 1 >= sql.size()) {
				// No room for the final dollar
				break;
			}
			// Skip to the next dollar symbol
			idx_t start = i + 1;
			idx_t end = start;
			while (end < sql.size() && sql[end] != '$') {
				end++;
			}
			if (end >= sql.size()) {
				// No final dollar, continue as normal
				break;
			}
			if (end - start != dollar_quote_marker.size()) {
				// Length mismatch, cannot match
				break;
			}
			if (sql.compare(start, dollar_quote_marker.size(), dollar_quote_marker) != 0) {
				// marker mismatch
				break;
			}
			// Marker found! Revert to standard state
			size_t full_marker_len = dollar_quote_marker.size() + 2;
			string quoted = sql.substr(last_pos, (start + dollar_quote_marker.size() + 1) - last_pos);
			string content = quoted.substr(full_marker_len, quoted.size() - 2 * full_marker_len);
			content = StringUtil::Replace(content, "'", "''");
			quoted = "'" + content + "'";
			tokens.emplace_back(quoted, full_marker_len);
			dollar_quote_marker = string();
			state = TokenizeState::STANDARD;
			i = end;
			last_pos = i + 1;
			break;
		}
		default:
			throw InternalException("unrecognized tokenize state");
		}
	}

	// finished processing - check the final state
	switch (state) {
	case TokenizeState::STRING_LITERAL:
		last_pos++;
		break;
	case TokenizeState::SINGLE_LINE_COMMENT:
	case TokenizeState::MULTI_LINE_COMMENT:
	case TokenizeState::DOLLAR_QUOTED_STRING:
		// no suggestions in comments or dollar-quoted strings
		return false;
	default:
		break;
	}
	string last_word = sql.substr(last_pos, sql.size() - last_pos);
	OnLastToken(state, std::move(last_word), last_pos);
	return true;
}

void BaseTokenizer::OnStatementEnd(idx_t pos) {
	tokens.clear();
}

} // namespace duckdb
