#include "./token_collector.hpp"
#include "./tree_sitter_utils.hpp"

namespace lsp::language::ast
{
    // ===== TokenCollector 实现 =====

    TokenCollector::TokenCollector(const TokenCollectionOptions& options) :
        options_(options)
    {
    }

    std::vector<TokenInfo> TokenCollector::Collect(TSNode root, const std::string& source)
    {
        std::vector<TokenInfo> tokens;
        TSNode null_parent = { 0 };
        CollectRecursive(root, source, tokens, 0, null_parent);
        return tokens;
    }

    std::vector<TokenInfo> TokenCollector::CollectByType(
        TSNode root,
        const std::string& source,
        const std::vector<std::string>& types)
    {
        // 临时修改选项以仅包含指定类型
        TokenCollectionOptions original_options = options_;
        options_.include_types = std::unordered_set<std::string>(types.begin(), types.end());

        auto result = Collect(root, source);

        // 恢复原始选项
        options_ = original_options;

        return result;
    }

    std::vector<TokenInfo> TokenCollector::CollectInRange(
        TSNode root,
        const std::string& source,
        uint32_t start_line,
        uint32_t end_line)
    {
        // 临时修改选项以包含位置过滤
        TokenCollectionOptions original_options = options_;
        options_.has_location_filter = true;
        options_.filter_start_line = start_line;
        options_.filter_end_line = end_line;

        auto result = Collect(root, source);

        // 恢复原始选项
        options_ = original_options;

        return result;
    }

    std::vector<TokenInfo> TokenCollector::CollectLeafNodes(TSNode root, const std::string& source)
    {
        auto all_tokens = Collect(root, source);
        std::vector<TokenInfo> leaf_tokens;

        for (const auto& token : all_tokens)
        {
            // 叶子节点的特征：文本长度等于字节范围
            // 或者可以通过重新检查 TSNode 来确定
            leaf_tokens.push_back(token);
        }

        return leaf_tokens;
    }

    std::vector<TokenInfo> TokenCollector::FindTokensAtPosition(TSNode root, const std::string& source, uint32_t line, uint32_t column)
    {
        auto all_tokens = Collect(root, source);
        std::vector<TokenInfo> result;

        for (const auto& token : all_tokens)
        {
            const auto& loc = token.location;
            // 检查位置是否在 token 范围内
            if (loc.start_line <= line && line <= loc.end_line)
            {
                if (line == loc.start_line && line == loc.end_line)
                {
                    if (loc.start_column <= column && column <= loc.end_column)
                        result.push_back(token);
                }
                else if (line == loc.start_line)
                {
                    if (column >= loc.start_column)
                        result.push_back(token);
                }
                else if (line == loc.end_line)
                {
                    if (column <= loc.end_column)
                        result.push_back(token);
                }
                else
                {
                    result.push_back(token);
                }
            }
        }

        return result;
    }

    std::vector<TokenInfo> TokenCollector::FindTokensByText(TSNode root, const std::string& source, const std::string& text, bool exact_match)
    {
        auto all_tokens = Collect(root, source);
        std::vector<TokenInfo> result;

        for (const auto& token : all_tokens)
        {
            if (exact_match)
            {
                if (token.text == text)
                    result.push_back(token);
            }
            else
            {
                if (token.text.find(text) != std::string::npos)
                    result.push_back(token);
            }
        }

        return result;
    }

    std::vector<TokenInfo> TokenCollector::FindTokensBy(TSNode root, const std::string& source, std::function<bool(const TokenInfo&)> predicate)
    {
        auto all_tokens = Collect(root, source);
        std::vector<TokenInfo> result;

        for (const auto& token : all_tokens)
        {
            if (predicate(token))
                result.push_back(token);
        }

        return result;
    }

    uint32_t TokenCollector::CountTokensByType(TSNode root, const std::string& source, const std::string& type)
    {
        auto tokens = CollectByType(root, source, { type });
        return tokens.size();
    }

    std::vector<std::string> TokenCollector::GetUniqueTypes(TSNode root, const std::string& source)
    {
        auto tokens = Collect(root, source);
        std::unordered_set<std::string> unique_types;

        for (const auto& token : tokens)
        {
            unique_types.insert(token.type);
        }

        return std::vector<std::string>(unique_types.begin(), unique_types.end());
    }

    // ===== 私有方法实现 =====

    void TokenCollector::CollectRecursive(TSNode node, const std::string& source, std::vector<TokenInfo>& tokens, uint32_t depth, TSNode parent)
    {
        // 检查深度限制
        if (depth > options_.max_depth || depth < options_.min_depth)
            return;

        // 创建 TokenInfo
        TokenInfo info;
        FillTokenInfo(info, node, source, depth, parent);

        // 检查是否应该收集此节点
        if (ShouldCollectNode(info))
        {
            tokens.push_back(info);
        }

        // 递归处理子节点
        uint32_t child_count = ts_node_child_count(node);
        for (uint32_t i = 0; i < child_count; ++i)
        {
            TSNode child = ts_node_child(node, i);
            CollectRecursive(child, source, tokens, depth + 1, node);
        }
    }

    bool TokenCollector::ShouldCollectNode(const TokenInfo& info) const
    {
        // 检查命名节点过滤
        if (!options_.include_anonymous && !info.is_named)
            return false;

        // 检查注释过滤
        if (!options_.include_comments && IsCommentNode(info.type))
            return false;

        // 检查空白过滤
        if (!options_.include_whitespace && IsWhitespaceNode(info.type))
            return false;

        // 检查错误节点过滤
        if (!options_.include_errors && info.is_error)
            return false;

        // 检查缺失节点过滤
        if (!options_.include_missing && info.is_missing)
            return false;

        // 检查类型包含过滤
        if (!options_.include_types.empty())
        {
            if (options_.include_types.find(info.type) == options_.include_types.end())
                return false;
        }

        // 检查类型排除过滤
        if (!options_.exclude_types.empty())
        {
            if (options_.exclude_types.find(info.type) != options_.exclude_types.end())
                return false;
        }

        // 检查空文本过滤
        if (options_.skip_empty_text && info.text.empty())
            return false;

        // 检查文本长度过滤
        uint32_t text_len = info.text.length();
        if (text_len < options_.min_text_length || text_len > options_.max_text_length)
            return false;

        // 检查位置过滤
        if (options_.has_location_filter && !IsLocationInRange(info.location))
            return false;

        // 检查自定义过滤器
        if (options_.custom_filter && !options_.custom_filter(info))
            return false;

        return true;
    }

    void TokenCollector::FillTokenInfo(TokenInfo& info, TSNode node, const std::string& source, uint32_t depth, TSNode parent) const
    {
        // 基本信息
        info.location = ts::NodeLocation(node);
        info.type = ts_node_type(node);
        info.text = ts::Text(node, source);
        info.is_named = ts_node_is_named(node);
        info.depth = depth;

        // 错误和缺失标记
        info.is_error = ts_node_is_error(node);
        info.is_missing = ts_node_is_missing(node);

        // 父节点信息
        if (!ts_node_is_null(parent))
        {
            info.parent_type = ts_node_type(parent);

            // 查找当前节点在父节点中的索引
            uint32_t sibling_count = ts_node_child_count(parent);
            info.sibling_count = sibling_count;

            for (uint32_t i = 0; i < sibling_count; ++i)
            {
                TSNode sibling = ts_node_child(parent, i);
                if (ts_node_eq(sibling, node))
                {
                    info.child_index = i;
                    break;
                }
            }
        }
    }

    bool TokenCollector::IsCommentNode(const std::string& type) const
    {
        return type == "line_comment" ||
               type == "block_comment" ||
               type == "nested_comment" ||
               type == "comment";
    }

    bool TokenCollector::IsWhitespaceNode(const std::string& type) const
    {
        return type == "whitespace" ||
               type == " " ||
               type == "\n" ||
               type == "\r" ||
               type == "\t";
    }

    bool TokenCollector::IsLocationInRange(const Location& loc) const
    {
        if (!options_.has_location_filter)
            return true;

        // 检查范围是否有重叠
        if (loc.end_line < options_.filter_start_line)
            return false;
        if (loc.start_line > options_.filter_end_line)
            return false;

        return true;
    }

}
