#pragma once

#include <vector>
#include <optional>
#include <algorithm>
#include "./types.hpp"

namespace lsp::language::symbol
{
    class LocationIndex
    {
    public:
        struct LocationEntry
        {
            uint32_t start_byte;
            uint32_t end_byte;
            SymbolId symbol_id;

            bool operator<(const LocationEntry& other) const
            {
                if (start_byte != other.start_byte)
                    return start_byte < other.start_byte;
                return end_byte > other.end_byte;
            }
        };

        struct ScopeEntry
        {
            uint32_t start_byte;
            uint32_t end_byte;
            ScopeId scope_id;

            bool operator<(const ScopeEntry& other) const
            {
                if (start_byte != other.start_byte)
                    return start_byte < other.start_byte;
                return end_byte > other.end_byte;
            }
        };

        struct ReferenceEntry
        {
            uint32_t start_byte;
            uint32_t end_byte;
            SymbolId symbol_id;

            bool operator<(const ReferenceEntry& other) const
            {
                return start_byte < other.start_byte;
            }
        };

        void AddSymbol(SymbolId id, const ast::Location& location);
        void AddScope(ScopeId id, const ast::Location& range);
        void AddReference(SymbolId symbol_id, const ast::Location& location);
        void RebuildIndex();
        void Clear();

        std::optional<SymbolId> FindSymbolAt(const ast::Location& location) const;
        std::optional<ScopeId> FindScopeAt(const ast::Location& location) const;
        std::optional<SymbolId> FindReferenceAt(const ast::Location& location) const;

    private:
        // 符号定义位置索引(按起始位置排序)
        std::vector<LocationEntry> symbol_entries_;
        bool symbols_sorted_ = true;

        // 作用域位置索引(按起始位置排序)
        std::vector<ScopeEntry> scope_entries_;
        bool scopes_sorted_ = true;

        // 引用位置索引(位置 -> 符号ID)
        std::vector<ReferenceEntry> reference_entries_;
        bool references_sorted_ = true;

        // 辅助方法：在排序数组中查找位置（模板实现放在头文件）
        template<typename T, typename IdType>
        std::optional<IdType> FindInSortedEntries(const std::vector<T>& entries, const ast::Location& location, bool find_innermost = true) const
        {
            if (entries.empty())
                return std::nullopt;

            auto it = std::lower_bound(
                entries.begin(),
                entries.end(),
                location.start_byte,
                [](const T& entry, uint32_t pos) {
                    return entry.start_byte < pos;
                });

            if (it != entries.begin())
                --it;

            std::vector<IdType> candidates;
            for (; it != entries.end() && it->start_byte <= location.start_byte; ++it)
            {
                if (IsLocationInRange(location, it->start_byte, it->end_byte))
                {
                    if constexpr (std::is_same_v<T, LocationEntry> || std::is_same_v<T, ReferenceEntry>)
                        candidates.push_back(it->symbol_id);
                    else if constexpr (std::is_same_v<T, ScopeEntry>)
                        candidates.push_back(it->scope_id);
                }
            }

            if (candidates.empty())
                return std::nullopt;

            if (find_innermost)
                return candidates.back();
            else
                return candidates.front();
        }

        static bool IsLocationInRange(const ast::Location& location, uint32_t start, uint32_t end)
        {
            return location.start_byte >= start && location.start_byte < end;
        }
    };

}
