#include <gtest/gtest.h>

#include <fmt/format.h>

#include <dce-core/dsl/dsl.hpp>
#include <ranges>

using namespace dce::dsl;

namespace
{

struct ExpectedToken
{
	std::string_view str;
	/// If this is null, then we just want to "inject" a string, for example whitespace
	/// into the stream. It's kinda ugly doing it that way but it doens't matter here.
	std::optional<TokenKind> kind = {};
};

/// Generates a full list of expected tokens from a simplified representation.
/// Takes care of the end token for us.
auto gen_expected_tokens(std::vector<ExpectedToken> expected_tokens) -> std::vector<Token>
{
	// Using space, because intializing `std::string` from `"\0"` will create an empty string.
	expected_tokens.emplace_back(" ", TokenKind::End);

	auto line = 1z;
	auto column = 1z;
	auto offset = 0z;
	auto result = std::vector<Token>();
	result.reserve(expected_tokens.size());
	for(auto&& expected_token : std::move(expected_tokens)) {
		auto const len = expected_token.str.size();

		if(expected_token.kind) {
			result.emplace_back(
				SourceLoc {
					.line = line,
					.column = column,
					.offset = offset,
				},
				*expected_token.kind,
				static_cast<u8>(len)
			);
		} else if(expected_token.str.contains('\n')) {
			line++;
			offset++;
			column = 1;
			continue;
		}

		column += len;
		offset += len;
	}

	return result;
}

} // namespace

TEST(Tokenizer, TokenizesEmptySource)
{
	auto tokenizer = Tokenizer::from_source("");
	auto result = tokenizer.tokenize();
	ASSERT_TRUE(result.has_value());
	ASSERT_EQ(result->tokens.size(), 1);

	auto const& tok = result->tokens.at(0);
	ASSERT_EQ(tok.kind, TokenKind::End);
	ASSERT_EQ(tok.length, 1);
	ASSERT_EQ(tok.loc.line, 1);
	ASSERT_EQ(tok.loc.column, 1);
	ASSERT_EQ(tok.loc.offset, 0);
}

TEST(Tokenizer, TokenizesAtoms)
{
	static auto const tokens = std::unordered_map<std::string_view, TokenKind> {
		{ "*", TokenKind::Star },
		{ "@", TokenKind::At },
		{ "$", TokenKind::Dollar },
		{ "[", TokenKind::LBracket },
		{ "]", TokenKind::RBracket },
		{ ",", TokenKind::Comma },
	};

	for(auto const& [c, t] : tokens) {
		auto tokenizer = Tokenizer::from_source(c);
		auto result = tokenizer.tokenize();
		ASSERT_TRUE(result.has_value());
		ASSERT_EQ(result->tokens.size(), 2);

		auto const& tok = result->tokens.at(0);
		ASSERT_EQ(tok.kind, t);
		ASSERT_EQ(tok.length, 1);
		ASSERT_EQ(tok.loc.line, 1);
		ASSERT_EQ(tok.loc.column, 1);
		ASSERT_EQ(tok.loc.offset, 0);

		auto const& end = result->tokens.at(1);
		ASSERT_EQ(end.kind, TokenKind::End);
		ASSERT_EQ(end.length, 1);
		ASSERT_EQ(end.loc.line, 1);
		ASSERT_EQ(end.loc.column, 2);
		ASSERT_EQ(end.loc.offset, 1);
	}
}

TEST(Tokenizer, DisallowsNullBytes)
{
	{
		static constexpr auto& nulls = "hello\0world";

		auto tokenizer = Tokenizer::from_source(std::string_view(rg::data(nulls), rg::size(nulls)));
		auto result = tokenizer.tokenize();
		ASSERT_FALSE(result.has_value());
		ASSERT_TRUE(result.error().is<TE_NullByte>());

		auto const nb = result.error().as<TE_NullByte>();
		ASSERT_EQ(nb->loc.line, 1);
		ASSERT_EQ(nb->loc.offset, 5);
		ASSERT_EQ(nb->loc.column, nb->loc.offset + 1);
	}
}

TEST(Tokenizer, DisallowsNullBytesManyLines)
{
	{
		/*
		  012345
		1 some text
		2 world
		3 hello\0world
		  012345
		*/
		static constexpr auto& nulls = "some text\nworld\nhello\0world";
		auto tokenizer = Tokenizer::from_source(std::string_view(rg::data(nulls), rg::size(nulls)));
		auto result = tokenizer.tokenize();
		ASSERT_FALSE(result.has_value());
		ASSERT_TRUE(result.error().is<TE_NullByte>());

		auto const nb = result.error().as<TE_NullByte>();
		ASSERT_EQ(nb->loc.offset, 21);
		ASSERT_EQ(nb->loc.line, 3);
		ASSERT_EQ(nb->loc.column, 6);
	}
}

TEST(Tokenizer, TokenizesSingleGuild)
{
	{
		auto result = Tokenizer::from_source("$guild"sv).tokenize();
		auto const expected = gen_expected_tokens({
			{ "$", TokenKind::Dollar },
			{ "guild", TokenKind::Identifier },
		});

		ASSERT_TRUE(result.has_value());
		ASSERT_EQ(result.value().tokens, expected);
	}
}

TEST(Tokenizer, TokenizesGuildWithSingleChannel)
{
	{
		auto result = Tokenizer::from_source("$guild[channel]"sv).tokenize();
		auto const expected = gen_expected_tokens({
			{ "$", TokenKind::Dollar },
			{ "guild", TokenKind::Identifier },
			{ "[", TokenKind::LBracket },
			{ "channel", TokenKind::Identifier },
			{ "]", TokenKind::RBracket },
		});

		ASSERT_TRUE(result.has_value());
		ASSERT_EQ(result.value().tokens, expected);
	}
	{
		auto result = Tokenizer::from_source("$1239306837317411816[716384436157349924,]"sv).tokenize();
		auto const expected = gen_expected_tokens({
			{ "$", TokenKind::Dollar },
			{ "1239306837317411816", TokenKind::Identifier },
			{ "[", TokenKind::LBracket },
			{ "716384436157349924", TokenKind::Identifier },
			{ ",", TokenKind::Comma },
			{ "]", TokenKind::RBracket },
		});

		ASSERT_TRUE(result.has_value());
		ASSERT_EQ(result.value().tokens, expected);
	}
}

TEST(Tokenizer, TokenizesGuildWithManyChannels)
{
	{
		auto result = Tokenizer::from_source("$guild[ch1, 716384436157349924, ch2, ch3,ch4,]"sv).tokenize();
		auto const expected = gen_expected_tokens({
			{ "$", TokenKind::Dollar },
			{ "guild", TokenKind::Identifier },
			{ "[", TokenKind::LBracket },
			{ "ch1", TokenKind::Identifier },
			{ ",", TokenKind::Comma },
			{ " " },
			{ "716384436157349924", TokenKind::Identifier },
			{ ",", TokenKind::Comma },
			{ " " },
			{ "ch2", TokenKind::Identifier },
			{ ",", TokenKind::Comma },
			{ " " },
			{ "ch3", TokenKind::Identifier },
			{ ",", TokenKind::Comma },
			{ "ch4", TokenKind::Identifier },
			{ ",", TokenKind::Comma },
			{ "]", TokenKind::RBracket },
		});

		ASSERT_TRUE(result.has_value());
		ASSERT_EQ(result.value().tokens, expected);
	}
}

TEST(Tokenizer, TokenizerHandlesWhitespace)
{
	{
		auto result = Tokenizer::from_source("  $guild [ch1, 716384436157349924  , ch2, ch3,ch4,]"sv).tokenize();
		auto const expected = gen_expected_tokens({
			{ "  " },
			{ "$", TokenKind::Dollar },
			{ "guild", TokenKind::Identifier },
			{ " " },
			{ "[", TokenKind::LBracket },
			{ "ch1", TokenKind::Identifier },
			{ ",", TokenKind::Comma },
			{ " " },
			{ "716384436157349924", TokenKind::Identifier },
			{ "  " },
			{ ",", TokenKind::Comma },
			{ " " },
			{ "ch2", TokenKind::Identifier },
			{ ",", TokenKind::Comma },
			{ " " },
			{ "ch3", TokenKind::Identifier },
			{ ",", TokenKind::Comma },
			{ "ch4", TokenKind::Identifier },
			{ ",", TokenKind::Comma },
			{ "]", TokenKind::RBracket },
		});

		ASSERT_TRUE(result.has_value());
		ASSERT_EQ(result.value().tokens, expected);
	}
}

TEST(Tokenizer, TokenizerHandlesNewlines)
{
	{
		auto result = Tokenizer::from_source("$guild[\nch1,\nch2,ch3\n]"sv).tokenize();
		auto const expected = gen_expected_tokens({
			{ "$", TokenKind::Dollar },
			{ "guild", TokenKind::Identifier },
			{ "[", TokenKind::LBracket },
			{ "\n" },
			{ "ch1", TokenKind::Identifier },
			{ ",", TokenKind::Comma },
			{ "\n" },
			{ "ch2", TokenKind::Identifier },
			{ ",", TokenKind::Comma },
			{ "ch3", TokenKind::Identifier },
			{ "\n" },
			{ "]", TokenKind::RBracket },
		});

		ASSERT_TRUE(result.has_value());
		ASSERT_EQ(result.value().tokens, expected);
	}
}

TEST(Tokenizer, TokenizesWholeSelection)
{
	{
		auto result = Tokenizer::from_source("@, $guild,\n\n$guild2[channel1, channel2,]\n"sv).tokenize();
		auto const expected = gen_expected_tokens({
			{ "@", TokenKind::At },
			{ ",", TokenKind::Comma },
			{ " " },
			{ "$", TokenKind::Dollar },
			{ "guild", TokenKind::Identifier },
			{ ",", TokenKind::Comma },
			{ "\n" },
			{ "\n" },
			{ "$", TokenKind::Dollar },
			{ "guild2", TokenKind::Identifier },
			{ "[", TokenKind::LBracket },
			{ "channel1", TokenKind::Identifier },
			{ ",", TokenKind::Comma },
			{ " " },
			{ "channel2", TokenKind::Identifier },
			{ ",", TokenKind::Comma },
			{ "]", TokenKind::RBracket },
			{ "\n" },
		});

		ASSERT_TRUE(result.has_value());
		ASSERT_EQ(result.value().tokens, expected);
	}
}