#include <cassert>
#include <cstdint>
#include <initializer_list>
#include <optional>
#include <vector>

#include "test/test.h"
#include "test/tester.h"
#include "feature/feature.h"
#include "semantic/symbol_kind.h"

#include "kota/meta/enum.h"

namespace clice::testing {

namespace {

namespace protocol = kota::ipc::protocol;

struct DecodedToken {
    LocalSourceRange range;
    std::uint32_t line = 0;
    std::uint32_t start = 0;
    std::uint32_t length = 0;
    std::uint32_t type = 0;
    std::uint32_t modifiers = 0;
};

auto compute_line_starts(llvm::StringRef content) -> std::vector<std::uint32_t> {
    std::vector<std::uint32_t> starts = {0};
    for(std::uint32_t i = 0; i < content.size(); ++i) {
        if(content[i] == '\n') {
            starts.push_back(i + 1);
        }
    }
    return starts;
}

auto decode_utf8_tokens(llvm::StringRef content, const protocol::SemanticTokens& tokens)
    -> std::vector<DecodedToken> {
    assert(tokens.data.size() % 5 == 0 && "invalid semantic token payload");

    auto starts = compute_line_starts(content);
    std::vector<DecodedToken> result;
    result.reserve(tokens.data.size() / 5);

    std::uint32_t line = 0;
    std::uint32_t character = 0;
    for(std::size_t i = 0; i < tokens.data.size(); i += 5) {
        auto delta_line = tokens.data[i + 0];
        auto delta_char = tokens.data[i + 1];
        auto length = tokens.data[i + 2];
        auto type = tokens.data[i + 3];
        auto modifiers = tokens.data[i + 4];

        line += delta_line;
        character = delta_line == 0 ? character + delta_char : delta_char;
        assert(line < starts.size() && "line index out of range");

        auto begin = starts[line] + character;
        auto end = begin + length;
        result.push_back({
            .range = LocalSourceRange(begin, end),
            .line = line,
            .start = character,
            .length = length,
            .type = type,
            .modifiers = modifiers,
        });
    }

    return result;
}

auto decode_relative_tokens(const protocol::SemanticTokens& tokens) -> std::vector<DecodedToken> {
    assert(tokens.data.size() % 5 == 0 && "invalid semantic token payload");

    std::vector<DecodedToken> result;
    result.reserve(tokens.data.size() / 5);

    std::uint32_t line = 0;
    std::uint32_t character = 0;
    for(std::size_t i = 0; i < tokens.data.size(); i += 5) {
        auto delta_line = tokens.data[i + 0];
        auto delta_char = tokens.data[i + 1];
        auto length = tokens.data[i + 2];
        auto type = tokens.data[i + 3];
        auto modifiers = tokens.data[i + 4];

        line += delta_line;
        character = delta_line == 0 ? character + delta_char : delta_char;
        result.push_back({
            .line = line,
            .start = character,
            .length = length,
            .type = type,
            .modifiers = modifiers,
        });
    }

    return result;
}

TEST_SUITE(semantic_tokens, Tester) {

protocol::SemanticTokens tokens;
std::vector<DecodedToken> decoded;

auto modifier_mask(std::initializer_list<SymbolModifiers::Kind> kinds) -> std::uint32_t {
    std::uint32_t mask = 0;
    for(auto kind: kinds) {
        mask |= SymbolModifiers::to_mask(kind);
    }
    return mask;
}

void run_utf8(llvm::StringRef code) {
    add_main("main.cpp", code);
    ASSERT_TRUE(compile_with_pch());
    tokens = feature::semantic_tokens(*unit, feature::PositionEncoding::UTF8);
    decoded = decode_utf8_tokens(unit->interested_content(), tokens);
}

auto find_by_range(llvm::StringRef name) -> const DecodedToken* {
    auto expected = range(name);
    for(const auto& token: decoded) {
        if(token.range == expected) {
            return &token;
        }
    }
    return nullptr;
}

void EXPECT_TOKEN(llvm::StringRef name,
                  SymbolKind::Kind expected_kind,
                  std::uint32_t expected_modifiers = 0) {
    auto* token = find_by_range(name);
    ASSERT_TRUE(token != nullptr);
    ASSERT_EQ(token->type, static_cast<std::uint32_t>(expected_kind));
    ASSERT_EQ(token->modifiers, expected_modifiers);
}

void EXPECT_NO_TOKEN(llvm::StringRef name) {
    ASSERT_TRUE(find_by_range(name) == nullptr);
}

TEST_CASE(BasicLexicalKinds) {
    run_utf8(R"cpp(
@d1[#define] @m0[FOO]
@k0[int] main() { @k1[return] 0; }
@c0[// comment]
)cpp");

    EXPECT_TOKEN("d1", SymbolKind::Directive);
    EXPECT_TOKEN("m0", SymbolKind::Macro);
    EXPECT_TOKEN("k0", SymbolKind::Keyword);
    EXPECT_TOKEN("k1", SymbolKind::Keyword);
    EXPECT_TOKEN("c0", SymbolKind::Comment);
}

TEST_CASE(IncludeDirective) {
    add_file("fake.h", "// fake header\n");
    add_main("main.cpp", R"cpp(
@d0[#include] @h0["fake.h"]
int main() { return 0; }
)cpp");
    ASSERT_TRUE(compile_with_pch());
    tokens = feature::semantic_tokens(*unit, feature::PositionEncoding::UTF8);
    decoded = decode_utf8_tokens(unit->interested_content(), tokens);

    EXPECT_TOKEN("d0", SymbolKind::Directive);
    EXPECT_TOKEN("h0", SymbolKind::Header);
}

TEST_CASE(LegacyIncludeForms) {
    add_file("fake.h", "// fake header\n");
    add_main("main.cpp", R"cpp(
@i0[#include] @h0["fake.h"]
@i1[#include] @h1["fake.h"]
@i2[#] @i3[include] @h2["fake.h"]
)cpp");
    ASSERT_TRUE(compile_with_pch());
    tokens = feature::semantic_tokens(*unit, feature::PositionEncoding::UTF8);
    decoded = decode_utf8_tokens(unit->interested_content(), tokens);

    EXPECT_TOKEN("i0", SymbolKind::Directive);
    EXPECT_TOKEN("h0", SymbolKind::Header);
    EXPECT_TOKEN("i1", SymbolKind::Directive);
    EXPECT_TOKEN("h1", SymbolKind::Header);
    EXPECT_TOKEN("i2", SymbolKind::Directive);
    EXPECT_TOKEN("i3", SymbolKind::Directive);
    EXPECT_TOKEN("h2", SymbolKind::Header);
}

TEST_CASE(LegacyComment) {
    run_utf8(R"cpp(
@line[/// line comment]
int x = 1;
)cpp");

    EXPECT_TOKEN("line", SymbolKind::Comment);
}

TEST_CASE(LegacyKeyword) {
    run_utf8(R"cpp(
@k0[int] main() {
    @k1[return] 0;
}
)cpp");

    EXPECT_TOKEN("k0", SymbolKind::Keyword);
    EXPECT_TOKEN("k1", SymbolKind::Keyword);
}

TEST_CASE(LegacyMacro) {
    run_utf8(R"cpp(
@directive[#define] @macro[FOO]
)cpp");

    EXPECT_TOKEN("directive", SymbolKind::Directive);
    EXPECT_TOKEN("macro", SymbolKind::Macro);
}

TEST_CASE(LegacyFinalAndOverride) {
    run_utf8(R"cpp(
struct A @final[final] {};

struct B {
    virtual void foo();
};

struct C : B {
    void foo() @override[override];
};

struct D : C {
    void foo() @final2[final];
};
)cpp");

    EXPECT_TOKEN("final", SymbolKind::Keyword);
    EXPECT_TOKEN("override", SymbolKind::Keyword);
    EXPECT_TOKEN("final2", SymbolKind::Keyword);
}

TEST_CASE(DeclarationAndTemplateModifiers) {
    run_utf8(R"cpp(
extern int @x1[x];
int @x2[x] = 0;

template <typename T>
extern int @y1[y];

template <typename T>
int @y2[y] = 0;

int main() {
    @x3[x] = 1;
}
)cpp");

    auto declaration = modifier_mask({SymbolModifiers::Declaration});
    auto definition = modifier_mask({SymbolModifiers::Definition});
    auto templated = modifier_mask({SymbolModifiers::Templated});

    EXPECT_TOKEN("x1", SymbolKind::Variable, declaration);
    EXPECT_TOKEN("x2", SymbolKind::Variable, definition);
    EXPECT_TOKEN("y1", SymbolKind::Variable, declaration | templated);
    EXPECT_TOKEN("y2", SymbolKind::Variable, definition | templated);
    EXPECT_TOKEN("x3", SymbolKind::Variable, 0);
}

TEST_CASE(IneligibleOperatorReferenceIsSuppressed) {
    run_utf8(R"cpp(
struct S {};

S operator+(S lhs, S rhs);

void use(S lhs, S rhs) {
    (void)(lhs @plus[+] rhs);
}
)cpp");

    EXPECT_NO_TOKEN("plus");
}

TEST_CASE(ConstructorAndDestructorNamesRemainHighlighted) {
    run_utf8(R"cpp(
struct S {
    @ctor_decl[S]();
    @dtor_decl[~]S();
};

S::@ctor_def[S]() {}

void use(S* value) {
    value->@dtor_ref[~]S();
}
)cpp");

    auto declaration = modifier_mask({SymbolModifiers::Declaration});
    auto definition = modifier_mask({SymbolModifiers::Definition});
    auto special_member = modifier_mask({SymbolModifiers::ConstructorOrDestructor});

    EXPECT_TOKEN("ctor_decl", SymbolKind::Method, declaration | special_member);
    EXPECT_TOKEN("dtor_decl", SymbolKind::Method, declaration | special_member);
    EXPECT_TOKEN("ctor_def", SymbolKind::Method, definition | special_member);
    EXPECT_TOKEN("dtor_ref", SymbolKind::Method, special_member);
}

TEST_CASE(LegacyVarDeclTemplates) {
    run_utf8(R"cpp(
extern int @x1[x];

int @x2[x] = 1;

template <typename T, typename U>
extern int @y1[y];

template <typename T, typename U>
int @y2[y] = 2;

template<typename T>
extern int @y3[y]<T, int>;

template<typename T>
int @y4[y]<T, int> = 4;

template<>
int @y5[y]<int, int> = 5;

int main() {
    @x3[x] = 6;
}
)cpp");

    auto declaration = modifier_mask({SymbolModifiers::Declaration});
    auto definition = modifier_mask({SymbolModifiers::Definition});
    auto templated = modifier_mask({SymbolModifiers::Templated});

    EXPECT_TOKEN("x1", SymbolKind::Variable, declaration);
    EXPECT_TOKEN("x2", SymbolKind::Variable, definition);
    EXPECT_TOKEN("y1", SymbolKind::Variable, declaration | templated);
    EXPECT_TOKEN("y2", SymbolKind::Variable, definition | templated);
    EXPECT_TOKEN("y3", SymbolKind::Variable, declaration | templated);
    EXPECT_TOKEN("y4", SymbolKind::Variable, definition | templated);
    EXPECT_TOKEN("y5", SymbolKind::Variable, definition);
    EXPECT_TOKEN("x3", SymbolKind::Variable, 0);
}

TEST_CASE(LegacyFunctionDecl) {
    run_utf8(R"cpp(
extern int @foo1[foo]();

int @foo2[foo]() {
    return 0;
}

template <typename T>
extern int @bar1[bar]();

template <typename T>
int @bar2[bar]() {
    return 1;
}
)cpp");

    auto declaration = modifier_mask({SymbolModifiers::Declaration});
    auto definition = modifier_mask({SymbolModifiers::Definition});
    auto templated = modifier_mask({SymbolModifiers::Templated});

    EXPECT_TOKEN("foo1", SymbolKind::Function, declaration);
    EXPECT_TOKEN("foo2", SymbolKind::Function, definition);
    EXPECT_TOKEN("bar1", SymbolKind::Function, declaration | templated);
    EXPECT_TOKEN("bar2", SymbolKind::Function, definition | templated);
}

TEST_CASE(LegacyRecordDecl) {
    run_utf8(R"cpp(
class @a1[A];

class @a2[A] {};

struct @b1[B];

struct @b2[B] {};

union @c1[C];

union @c2[C] {};
)cpp");

    auto declaration = modifier_mask({SymbolModifiers::Declaration});
    auto definition = modifier_mask({SymbolModifiers::Definition});

    EXPECT_TOKEN("a1", SymbolKind::Class, declaration);
    EXPECT_TOKEN("a2", SymbolKind::Class, definition);
    EXPECT_TOKEN("b1", SymbolKind::Struct, declaration);
    EXPECT_TOKEN("b2", SymbolKind::Struct, definition);
    EXPECT_TOKEN("c1", SymbolKind::Union, declaration);
    EXPECT_TOKEN("c2", SymbolKind::Union, definition);
}

TEST_CASE(UTF16LengthDiffersFromUTF8) {
    add_main("main.cpp", R"cpp(
int main() {
@lit[u8"你"];
}
)cpp");
    ASSERT_TRUE(compile_with_pch());

    auto utf8_tokens = feature::semantic_tokens(*unit, feature::PositionEncoding::UTF8);
    auto utf16_tokens = feature::semantic_tokens(*unit, feature::PositionEncoding::UTF16);

    auto utf8 = decode_utf8_tokens(unit->interested_content(), utf8_tokens);
    auto utf16 = decode_relative_tokens(utf16_tokens);

    auto string_type = static_cast<std::uint32_t>(SymbolKind::String);
    auto lit_range = range("lit");

    std::optional<DecodedToken> utf8_token;
    for(const auto& token: utf8) {
        if(token.range == lit_range && token.type == string_type) {
            utf8_token = token;
            break;
        }
    }
    ASSERT_TRUE(utf8_token.has_value());

    std::optional<DecodedToken> utf16_token;
    for(const auto& token: utf16) {
        if(token.line == utf8_token->line && token.start == utf8_token->start &&
           token.type == string_type) {
            utf16_token = token;
            break;
        }
    }
    ASSERT_TRUE(utf16_token.has_value());

    ASSERT_TRUE(utf8_token->length > utf16_token->length);
}

TEST_CASE(MultiLineCommentSplitMatchesLegacyConverter) {
    add_main("main.cpp", R"cpp(
int main() {
/*ab
cd*/
}
)cpp");
    ASSERT_TRUE(compile_with_pch());

    auto utf8_tokens = feature::semantic_tokens(*unit, feature::PositionEncoding::UTF8);
    auto relative = decode_relative_tokens(utf8_tokens);

    auto comment_type = static_cast<std::uint32_t>(SymbolKind::Comment);
    std::vector<DecodedToken> comments;
    for(const auto& token: relative) {
        if(token.type == comment_type) {
            comments.push_back(token);
        }
    }

    ASSERT_EQ(comments.size(), 2);
    ASSERT_EQ(comments[0].length, 5);
    ASSERT_EQ(comments[1].line, comments[0].line + 1);
    ASSERT_EQ(comments[1].start, 0);
    ASSERT_EQ(comments[1].length, 4);
}

TEST_CASE(ModuleDeclaration) {
    add_main("main.cpp", R"cpp(
export @kw[module] @mod[foo];
)cpp");
    ASSERT_TRUE(compile("-std=c++20"));
    tokens = feature::semantic_tokens(*unit, feature::PositionEncoding::UTF8);
    decoded = decode_utf8_tokens(unit->interested_content(), tokens);

    EXPECT_TOKEN("kw", SymbolKind::Keyword);
    EXPECT_TOKEN("mod", SymbolKind::Module);
}

TEST_CASE(ModuleDeclarationDotted) {
    add_main("main.cpp", R"cpp(
export @kw[module] @m0[foo].@m1[bar];
)cpp");
    ASSERT_TRUE(compile("-std=c++20"));
    tokens = feature::semantic_tokens(*unit, feature::PositionEncoding::UTF8);
    decoded = decode_utf8_tokens(unit->interested_content(), tokens);

    EXPECT_TOKEN("kw", SymbolKind::Keyword);
    EXPECT_TOKEN("m0", SymbolKind::Module);
    EXPECT_TOKEN("m1", SymbolKind::Module);
}

TEST_CASE(ModuleImport) {
    add_files("main.cpp", R"(
#[mod.cppm]
export module foo;
export int x = 42;

#[main.cpp]
@kw[import] @mod[foo];
int y = x;
)");
    ASSERT_TRUE(compile_with_modules());
    tokens = feature::semantic_tokens(*unit, feature::PositionEncoding::UTF8);
    decoded = decode_utf8_tokens(unit->interested_content(), tokens);

    EXPECT_TOKEN("kw", SymbolKind::Keyword);
    EXPECT_TOKEN("mod", SymbolKind::Module);
}

TEST_CASE(ModulePartition) {
    add_main("main.cpp", R"cpp(
export module @m0[foo]:@m1[bar];
)cpp");
    ASSERT_TRUE(compile("-std=c++20"));
    tokens = feature::semantic_tokens(*unit, feature::PositionEncoding::UTF8);
    decoded = decode_utf8_tokens(unit->interested_content(), tokens);

    EXPECT_TOKEN("m0", SymbolKind::Module);
    EXPECT_TOKEN("m1", SymbolKind::Module);
}

TEST_CASE(ModuleReexport) {
    add_files("main.cppm", R"(
#[mod.cppm]
export module foo;
export int x = 42;

#[main.cppm]
export module bar;
export @kw[import] @mod[foo];
)");
    ASSERT_TRUE(compile_with_modules());
    tokens = feature::semantic_tokens(*unit, feature::PositionEncoding::UTF8);
    decoded = decode_utf8_tokens(unit->interested_content(), tokens);

    EXPECT_TOKEN("kw", SymbolKind::Keyword);
    EXPECT_TOKEN("mod", SymbolKind::Module);
}

TEST_CASE(GlobalModuleFragment) {
    add_main("main.cpp", R"cpp(
module;
export module @mod[foo];
)cpp");
    ASSERT_TRUE(compile("-std=c++20"));
    tokens = feature::semantic_tokens(*unit, feature::PositionEncoding::UTF8);
    decoded = decode_utf8_tokens(unit->interested_content(), tokens);

    EXPECT_TOKEN("mod", SymbolKind::Module);
}

TEST_CASE(PrivateModuleFragment) {
    add_main("main.cpp", R"cpp(
export module @mod[foo];
module :private;
int x = 1;
)cpp");
    ASSERT_TRUE(compile("-std=c++20"));
    tokens = feature::semantic_tokens(*unit, feature::PositionEncoding::UTF8);
    decoded = decode_utf8_tokens(unit->interested_content(), tokens);

    EXPECT_TOKEN("mod", SymbolKind::Module);
}

TEST_CASE(ModuleKeywordAsIdentifier) {
    run_utf8(R"cpp(
void f() {
    struct @s0[module] {};
    @s1[module] @v0[m];
    int @v1[import] = 1;
    int @v2[module] = 2;
}
)cpp");

    auto definition = modifier_mask({SymbolModifiers::Definition});
    EXPECT_TOKEN("s0", SymbolKind::Struct, definition);
    EXPECT_TOKEN("s1", SymbolKind::Struct);
    EXPECT_TOKEN("v0", SymbolKind::Variable, definition);
    EXPECT_TOKEN("v1", SymbolKind::Variable, definition);
    EXPECT_TOKEN("v2", SymbolKind::Variable, definition);
}

TEST_CASE(snapshot) {
    ASSERT_SNAPSHOT_GLOB(corpus_dir, "**/*.cpp", [&](std::string_view path) -> std::string {
        if(!compile_file(path))
            return "COMPILE_ERROR";
        auto content = unit->interested_content();
        auto tokens = feature::semantic_tokens(*unit);
        feature::PositionMapper mapper(content, feature::PositionEncoding::UTF8);
        std::string result;
        for(auto& token: tokens) {
            if(!token.range.valid() || token.range.end <= token.range.begin ||
               token.range.end > content.size())
                continue;

            auto pos = mapper.to_position(token.range.begin);
            if(!pos)
                continue;

            auto text = content.substr(token.range.begin, token.range.length());
            auto kind = kota::meta::enum_name(static_cast<SymbolKind::Kind>(token.kind), "Unknown");

            result += std::format("- {{ loc: \"{}:{}\", text: {}, kind: {}",
                                  pos->line,
                                  pos->character,
                                  yaml_str(text),
                                  kind);

            std::string mods;
            for(std::uint32_t i = 0; i < 32; ++i) {
                if(token.modifiers & (1u << i)) {
                    auto name = kota::meta::enum_name(static_cast<SymbolModifiers::Kind>(i));
                    if(!name.empty()) {
                        if(!mods.empty())
                            mods += ", ";
                        mods += name;
                    }
                }
            }
            if(!mods.empty()) {
                result += std::format(", modifiers: [{}]", mods);
            }
            result += " }\n";
        }
        return result;
    });
}

};  // TEST_SUITE(semantic_tokens)

}  // namespace

}  // namespace clice::testing