#include "CppHighlighter.h"
#include "../Buffer.h"
#include <cctype>

namespace kte {
static bool
is_digit(char c)
{
	return c >= '0' && c <= '9';
}


CppHighlighter::CppHighlighter()
{
	const char *kw[] = {
		"if", "else", "for", "while", "do", "switch", "case", "default", "break", "continue",
		"return", "goto", "struct", "class", "namespace", "using", "template", "typename",
		"public", "private", "protected", "virtual", "override", "const", "constexpr", "auto",
		"static", "inline", "operator", "new", "delete", "try", "catch", "throw", "friend",
		"enum", "union", "extern", "volatile", "mutable", "noexcept", "sizeof", "this"
	};
	for (auto s: kw)
		keywords_.insert(s);
	const char *types[] = {
		"int", "long", "short", "char", "signed", "unsigned", "float", "double", "void",
		"bool", "wchar_t", "size_t", "ptrdiff_t", "uint8_t", "uint16_t", "uint32_t", "uint64_t",
		"int8_t", "int16_t", "int32_t", "int64_t"
	};
	for (auto s: types)
		types_.insert(s);
}


bool
CppHighlighter::is_ident_start(char c)
{
	return std::isalpha(static_cast<unsigned char>(c)) || c == '_';
}


bool
CppHighlighter::is_ident_char(char c)
{
	return std::isalnum(static_cast<unsigned char>(c)) || c == '_';
}


void
CppHighlighter::HighlightLine(const Buffer &buf, int row, std::vector<HighlightSpan> &out) const
{
	// Stateless entry simply delegates to stateful with a clean previous state
	StatefulHighlighter::LineState prev;
	(void) HighlightLineStateful(buf, row, prev, out);
}


StatefulHighlighter::LineState
CppHighlighter::HighlightLineStateful(const Buffer &buf,
                                      int row,
                                      const LineState &prev,
                                      std::vector<HighlightSpan> &out) const
{
	const auto &rows                     = buf.Rows();
	StatefulHighlighter::LineState state = prev;
	if (row < 0 || static_cast<std::size_t>(row) >= rows.size())
		return state;
	std::string s = static_cast<std::string>(rows[static_cast<std::size_t>(row)]);
	if (s.empty())
		return state;

	auto push = [&](int a, int b, TokenKind k) {
		if (b > a)
			out.push_back({a, b, k});
	};
	int n   = static_cast<int>(s.size());
	int bol = 0;
	while (bol < n && (s[bol] == ' ' || s[bol] == '\t'))
		++bol;
	int i = 0;

	// Continue multi-line raw string from previous line
	if (state.in_raw_string) {
		std::string needle = ")" + state.raw_delim + "\"";
		auto pos           = s.find(needle);
		if (pos == std::string::npos) {
			push(0, n, TokenKind::String);
			state.in_raw_string = true;
			return state;
		} else {
			int end = static_cast<int>(pos + needle.size());
			push(0, end, TokenKind::String);
			i                   = end;
			state.in_raw_string = false;
			state.raw_delim.clear();
		}
	}

	// Continue multi-line block comment from previous line
	if (state.in_block_comment) {
		int j = i;
		while (i + 1 < n) {
			if (s[i] == '*' && s[i + 1] == '/') {
				i += 2;
				push(j, i, TokenKind::Comment);
				state.in_block_comment = false;
				break;
			}
			++i;
		}
		if (state.in_block_comment) {
			push(j, n, TokenKind::Comment);
			return state;
		}
	}

	while (i < n) {
		char c = s[i];
		// Preprocessor at beginning of line (after leading whitespace)
		if (i == bol && c == '#') {
			push(0, n, TokenKind::Preproc);
			break;
		}

		// Whitespace
		if (c == ' ' || c == '\t') {
			int j = i + 1;
			while (j < n && (s[j] == ' ' || s[j] == '\t'))
				++j;
			push(i, j, TokenKind::Whitespace);
			i = j;
			continue;
		}

		// Line comment
		if (c == '/' && i + 1 < n && s[i + 1] == '/') {
			push(i, n, TokenKind::Comment);
			break;
		}

		// Block comment
		if (c == '/' && i + 1 < n && s[i + 1] == '*') {
			int j       = i + 2;
			bool closed = false;
			while (j + 1 <= n) {
				if (j + 1 < n && s[j] == '*' && s[j + 1] == '/') {
					j += 2;
					closed = true;
					break;
				}
				++j;
			}
			if (closed) {
				push(i, j, TokenKind::Comment);
				i = j;
				continue;
			}
			// Spill to next lines
			push(i, n, TokenKind::Comment);
			state.in_block_comment = true;
			return state;
		}

		// Raw string start: very simple detection: R"delim(
		if (c == 'R' && i + 1 < n && s[i + 1] == '"') {
			int k = i + 2;
			std::string delim;
			while (k < n && s[k] != '(') {
				delim.push_back(s[k]);
				++k;
			}
			if (k < n && s[k] == '(') {
				int body_start     = k + 1;
				std::string needle = ")" + delim + "\"";
				auto pos           = s.find(needle, static_cast<std::size_t>(body_start));
				if (pos == std::string::npos) {
					push(i, n, TokenKind::String);
					state.in_raw_string = true;
					state.raw_delim     = delim;
					return state;
				} else {
					int end = static_cast<int>(pos + needle.size());
					push(i, end, TokenKind::String);
					i = end;
					continue;
				}
			}
			// If malformed, just treat 'R' as identifier fallback
		}

		// Regular string literal
		if (c == '"') {
			int j    = i + 1;
			bool esc = false;
			while (j < n) {
				char d = s[j++];
				if (esc) {
					esc = false;
					continue;
				}
				if (d == '\\') {
					esc = true;
					continue;
				}
				if (d == '"')
					break;
			}
			push(i, j, TokenKind::String);
			i = j;
			continue;
		}

		// Char literal
		if (c == '\'') {
			int j    = i + 1;
			bool esc = false;
			while (j < n) {
				char d = s[j++];
				if (esc) {
					esc = false;
					continue;
				}
				if (d == '\\') {
					esc = true;
					continue;
				}
				if (d == '\'')
					break;
			}
			push(i, j, TokenKind::Char);
			i = j;
			continue;
		}

		// Number literal (simple)
		if (is_digit(c) || (c == '.' && i + 1 < n && is_digit(s[i + 1]))) {
			int j = i + 1;
			while (j < n && (std::isalnum(static_cast<unsigned char>(s[j])) || s[j] == '.' || s[j] == 'x' ||
			                 s[j] == 'X' || s[j] == 'b' || s[j] == 'B' || s[j] == '_'))
				++j;
			push(i, j, TokenKind::Number);
			i = j;
			continue;
		}

		// Identifier / keyword / type
		if (is_ident_start(c)) {
			int j = i + 1;
			while (j < n && is_ident_char(s[j]))
				++j;
			std::string id = s.substr(i, j - i);
			TokenKind k    = TokenKind::Identifier;
			if (keywords_.count(id))
				k = TokenKind::Keyword;
			else if (types_.count(id))
				k = TokenKind::Type;
			push(i, j, k);
			i = j;
			continue;
		}

		// Operators and punctuation (single char for now)
		TokenKind kind = TokenKind::Operator;
		if (std::ispunct(static_cast<unsigned char>(c)) && c != '_' && c != '#') {
			if (c == ';' || c == ',' || c == '(' || c == ')' || c == '{' || c == '}' || c == '[' || c ==
			    ']')
				kind = TokenKind::Punctuation;
			push(i, i + 1, kind);
			++i;
			continue;
		}

		// Fallback
		push(i, i + 1, TokenKind::Default);
		++i;
	}

	return state;
}
} // namespace kte