#include <cstring>
#include <mutex>
#include <unordered_map>
#include <string>
#include <tree_sitter/api.h>

#include "text/text.hpp"
#include "text/token.hpp"

namespace
Nimbin
{
extern "C" const TSLanguage* tree_sitter_cpp(void);

namespace {

thread_local TSParser* parser = [] {
    TSParser* p = ts_parser_new();
    ts_parser_set_language(p, tree_sitter_cpp());
    return p;
}();

// ---------------------------
// Keyword / literal fast path
// ---------------------------
Text::Token::Type classify_Keyword(const std::string& w)
{
    using T = Text::Token::Type;
    static const std::unordered_map<std::string, T> kw = {
        // ── Type / StorageClass / Structure / Typedef → Vim Type (SeaGreen) ──
        {"void",       T::Type}, {"bool",        T::Type}, {"char",         T::Type}, {"wchar_t",   T::Type},
        {"char8_t",    T::Type}, {"char16_t",    T::Type}, {"char32_t",     T::Type},
        {"short",      T::Type}, {"int",         T::Type}, {"long",         T::Type},
        {"float",      T::Type}, {"double",      T::Type}, {"signed",       T::Type}, {"unsigned",  T::Type},
        {"auto",       T::Type},
        {"size_t",     T::Type}, {"ssize_t",     T::Type}, {"ptrdiff_t",    T::Type},
        {"int8_t",     T::Type}, {"int16_t",     T::Type}, {"int32_t",      T::Type}, {"int64_t",   T::Type},
        {"uint8_t",    T::Type}, {"uint16_t",    T::Type}, {"uint32_t",     T::Type}, {"uint64_t",  T::Type},
        {"intptr_t",   T::Type}, {"uintptr_t",   T::Type},
        // StorageClass+ qualifiers (Vim links these to Type)
        {"const",      T::Type}, {"constexpr",   T::Type}, {"consteval",    T::Type}, {"constinit", T::Type},
        {"static",     T::Type}, {"extern",      T::Type}, {"inline",       T::Type}, {"mutable",   T::Type},
        {"volatile",   T::Type}, {"register",    T::Type}, {"thread_local", T::Type},
        // Structure + Typedef (Vim links these to Type)
        {"struct",     T::Type}, {"class",       T::Type}, {"enum",         T::Type}, {"union",     T::Type},
        {"typedef",    T::Type}, {"typename",    T::Type},

        // ── Constant (Magenta) ──────────────────────────────────────────
        {"true",       T::NumberLit}, {"false",  T::NumberLit},
        {"nullptr",    T::NumberLit}, {"NULL",   T::NumberLit},

        {"this",       T::Variable},   // Identifier (DarkCyan)

        // ── Statement/ Conditional / Repeat / Operator / Exception → Vim
        //    Statement(Brown). Also namespace/using/template (cppStatement). ──
        {"if",         T::Keyword}, {"else",     T::Keyword}, {"for",       T::Keyword}, {"while",  T::Keyword},
        {"do",         T::Keyword}, {"switch",   T::Keyword}, {"case",      T::Keyword}, {"default",T::Keyword},
        {"break",      T::Keyword}, {"continue", T::Keyword}, {"goto",      T::Keyword}, {"return", T::Keyword},
        {"new",        T::Keyword}, {"delete",   T::Keyword}, {"sizeof",    T::Keyword},
        {"alignof",    T::Keyword}, {"alignas",  T::Keyword}, {"typeid",    T::Keyword},
        {"operator",   T::Keyword}, {"decltype", T::Keyword}, {"requires",  T::Keyword},
        {"try",        T::Keyword}, {"catch",    T::Keyword}, {"throw",     T::Keyword},
        {"static_cast",T::Keyword}, {"dynamic_cast",    T::Keyword},
        {"const_cast", T::Keyword}, {"reinterpret_cast",T::Keyword},
        {"co_await",   T::Keyword}, {"co_yield", T::Keyword}, {"co_return", T::Keyword},
        {"namespace",  T::Keyword}, {"using",    T::Keyword}, {"template",  T::Keyword}, {"concept",T::Keyword},
        {"public",     T::Keyword}, {"private",  T::Keyword}, {"protected", T::Keyword},
        {"virtual",    T::Keyword}, {"override", T::Keyword}, {"final",     T::Keyword},
        {"explicit",   T::Keyword}, {"friend",   T::Keyword}, {"noexcept",  T::Keyword},
    };

    auto it = kw.find(w);
    if (it != kw.end())
        return it->second;

    return Text::Token::Type::Identifier;
}

// ---------------------------
// Tree-sitter node classifier
// ---------------------------
bool classify_Node(TSNode node, Text::Token::Type& out)
{
    const char* t = ts_node_type(node);

    // comments
    if (std::strcmp(t, "comment") == 0) {
        out = Text::Token::Type::Comment;
        return true;
    }

    // strings
    if (std::strcmp(t, "string_literal") == 0 ||
        std::strcmp(t, "raw_string_literal") == 0 ||
        std::strcmp(t, "char_literal") == 0 ||
        std::strcmp(t, "system_lib_string") == 0) {
        out = Text::Token::Type::StringLit;
        return true;
    }

    // numbers
    if (std::strcmp(t, "number_literal") == 0 ||
        std::strcmp(t, "integer_literal") == 0 ||
        std::strcmp(t, "floating_point_literal") == 0) {
        out = Text::Token::Type::NumberLit;
        return true;
    }

    // types
    if (std::strcmp(t, "primitive_type") == 0 ||
        std::strcmp(t, "type_identifier") == 0 ||
        std::strcmp(t, "sized_type_specifier") == 0 ||
        std::strcmp(t, "template_type") == 0) {
        out = Text::Token::Type::Type;
        return true;
    }

    // namespaces
    if (std::strcmp(t, "namespace_identifier") == 0) {
        out = Text::Token::Type::Namespace;
        return true;
    }

    // fields
    if (std::strcmp(t, "field_identifier") == 0) {
        out = Text::Token::Type::Variable;
        return true;
    }

    // operators / operator overloads
    if (std::strcmp(t, "operator_name") == 0) {
        out = Text::Token::Type::Function;
        return true;
    }

    // preprocessor (keep tight, not prefix-based)
    if (std::strncmp(t, "preproc_", 8) == 0) {
        out = Text::Token::Type::Preprocessor;
        return true;
    }

    // identifiers (context sensitive)
    if (std::strcmp(t, "identifier") == 0) {
        TSNode parent = ts_node_parent(node);
        const char* p = ts_node_is_null(parent) ? "" : ts_node_type(parent);

        if (std::strcmp(p, "function_declarator") == 0 ||
            std::strcmp(p, "call_expression") == 0 ||
            std::strcmp(p, "operator_name") == 0 ||
            std::strcmp(p, "template_function") == 0) {
            out = Text::Token::Type::Function;
            return true;
        }

        if (std::strcmp(p, "qualified_identifier") == 0 ||
            std::strcmp(p, "namespace_identifier") == 0) {
            out = Text::Token::Type::Namespace;
            return true;
        }

        out = Text::Token::Type::Variable;
        return true;
    }

    return false;
}

// ---------------------------
// Tree walk
// ---------------------------
void walk(TSNode node, DynArray<Text::HiSpan>& out)
{
    Text::Token::Type type;

    uint32_t count = ts_node_child_count(node);

    if (count == 0) {
        if (classify_Node(node, type)) {
            out.push_back({
                ts_node_start_byte(node),
                ts_node_end_byte(node),
                type
            });
        }
        return;
    }

    if (classify_Node(node, type)) {
        out.push_back({
            ts_node_start_byte(node),
            ts_node_end_byte(node),
            type
        });
        return;
    }

    for (uint32_t i = 0; i < count; i++)
        walk(ts_node_child(node, i), out);
}

} // namespace

// ---------------------------
// Public API
// ---------------------------
DynArray<Text::HiSpan>
Text::CppHighlight::classify(const std::string& source)
{
    DynArray<Text::HiSpan> spans;

    TSTree* tree = ts_parser_parse_string(
        parser,
        nullptr,
        source.c_str(),
        (uint32_t)source.size()
    );

    walk(ts_tree_root_node(tree), spans);

    ts_tree_delete(tree);
    return spans;
}

// ---------------------------
// Word classification (FIXED)
// ---------------------------
Text::Token::Type
Text::CppHighlight::classify_Word(const std::string& word)
{
    // All-whitespace (spaces/tab) → a whitespace token.
    if (!word.empty()) {
        bool only_ws = true;
        for (char c : word) if (c != ' ' && c != '\t') { only_ws = false; break; }
        if (only_ws) return Text::Token::Type::Whitespace;

        // Punctuation / operators / numeric literals — context-free, classify directly
        // so they don't fall through to Identifier when tree-sitter emits no span.
        const char c0 = word[0];
        if ((c0 >= '0' && c0 <= '9') ||
            (c0 == '.' && word.size() > 1 && word[1] >= '0' && word[1] <= '9'))
            return Text::Token::Type::NumberLit;

        auto all_of = [&](const char* set) {
            for (char c : word) {
                bool ok = false;
                for (const char* p = set; *p; ++p) if (*p == c) { ok = true; break; }
                if (!ok) return false;
            }
            return true;
        };
        if (all_of("()[]{};,"))            return Text::Token::Type::Delimiter;
        if (all_of("+-*/%=<>!&|^~.:?"))    return Text::Token::Type::Operator;
    }

    static std::unordered_map<std::string, Text::Token::Type> cache;
    static std::mutex mtx;

    {
        std::lock_guard<std::mutex> lk(mtx);
        auto it = cache.find(word);
        if (it != cache.end())
            return it->second;
    }

    // fast deterministic classification (NO tree-sitter here)
    Text::Token::Type type = classify_Keyword(word);

    {
        std::lock_guard<std::mutex> lk(mtx);
        cache[word] = type;
    }

    return type;
}

} // namespace Nimbin