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

namespace kte {
static void
push(std::vector<HighlightSpan> &out, int a, int b, TokenKind k)
{
	if (b > a)
		out.push_back({a, b, k});
}


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


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


PythonHighlighter::PythonHighlighter()
{
	const char *kw[] = {
		"and", "as", "assert", "break", "class", "continue", "def", "del", "elif", "else", "except", "False",
		"finally", "for", "from", "global", "if", "import", "in", "is", "lambda", "None", "nonlocal", "not",
		"or", "pass", "raise", "return", "True", "try", "while", "with", "yield"
	};
	for (auto s: kw)
		kws_.insert(s);
}


void
PythonHighlighter::HighlightLine(const Buffer &buf, int row, std::vector<HighlightSpan> &out) const
{
	LineState st;
	(void) HighlightLineStateful(buf, row, st, out);
}


StatefulHighlighter::LineState
PythonHighlighter::HighlightLineStateful(const Buffer &buf, int row, const LineState &prev,
                                         std::vector<HighlightSpan> &out) const
{
	StatefulHighlighter::LineState state = prev;
	const auto &rows                     = buf.Rows();
	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)]);
	int n         = static_cast<int>(s.size());

	// Triple-quoted string continuation uses in_raw_string with raw_delim either "'''" or "\"\"\""
	if (state.in_raw_string && (state.raw_delim == "'''" || state.raw_delim == "\"\"\"")) {
		auto pos = s.find(state.raw_delim);
		if (pos == std::string::npos) {
			push(out, 0, n, TokenKind::String);
			return state; // still inside
		} else {
			int end = static_cast<int>(pos + static_cast<int>(state.raw_delim.size()));
			push(out, 0, end, TokenKind::String);
			// remainder processed normally
			s                   = s.substr(end);
			n                   = static_cast<int>(s.size());
			state.in_raw_string = false;
			state.raw_delim.clear();
			// Continue parsing remainder as a separate small loop
			int base = end;
			// original offset, but we already emitted to 'out' with base=0; following spans should be from 'end'
			// For simplicity, mark rest as Default
			if (n > 0)
				push(out, base, base + n, TokenKind::Default);
			return state;
		}
	}

	int i = 0;
	// Detect comment start '#', ignoring inside strings
	while (i < n) {
		char c = s[i];
		if (c == ' ' || c == '\t') {
			int j = i + 1;
			while (j < n && (s[j] == ' ' || s[j] == '\t'))
				++j;
			push(out, i, j, TokenKind::Whitespace);
			i = j;
			continue;
		}
		if (c == '#') {
			push(out, i, n, TokenKind::Comment);
			break;
		}
		// Strings: triple quotes and single-line
		if (c == '"' || c == '\'') {
			char q = c;
			// triple?
			if (i + 2 < n && s[i + 1] == q && s[i + 2] == q) {
				std::string delim(3, q);
				int j    = i + 3; // search for closing triple
				auto pos = s.find(delim, static_cast<std::size_t>(j));
				if (pos == std::string::npos) {
					push(out, i, n, TokenKind::String);
					state.in_raw_string = true;
					state.raw_delim     = delim;
					return state;
				} else {
					int end = static_cast<int>(pos + 3);
					push(out, i, end, TokenKind::String);
					i = end;
					continue;
				}
			} else {
				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 == q)
						break;
				}
				push(out, i, j, TokenKind::String);
				i = j;
				continue;
			}
		}
		if (std::isdigit(static_cast<unsigned char>(c))) {
			int j = i + 1;
			while (j < n && (std::isalnum(static_cast<unsigned char>(s[j])) || s[j] == '.' || s[j] == '_'))
				++j;
			push(out, i, j, TokenKind::Number);
			i = j;
			continue;
		}
		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 (kws_.count(id))
				k = TokenKind::Keyword;
			push(out, i, j, k);
			i = j;
			continue;
		}
		if (std::ispunct(static_cast<unsigned char>(c))) {
			TokenKind k = TokenKind::Operator;
			if (c == ':' || c == ',' || c == '(' || c == ')' || c == '[' || c == ']')
				k = TokenKind::Punctuation;
			push(out, i, i + 1, k);
			++i;
			continue;
		}
		push(out, i, i + 1, TokenKind::Default);
		++i;
	}
	return state;
}
} // namespace kte