#pragma once

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

namespace lsp::language::symbol
{

    class LocationIndex : public ISymbolIndex
    {
    public:
        void OnSymbolAdded(const Symbol& symbol) override
        {
            const auto& range = symbol.range();
            entries_.push_back({ range.start_offset,
                                 range.end_offset,
                                 symbol.id() });
            needs_sort_ = true;
        }

        void OnSymbolRemoved(SymbolId id) override
        {
            entries_.erase(
                std::remove_if(entries_.begin(), entries_.end(), [id](const Entry& e) { return e.symbol_id == id; }),
                entries_.end());
        }

        void Clear() override
        {
            entries_.clear();
            needs_sort_ = false;
        }

        std::optional<SymbolId> FindAt(const ast::Location& location) const
        {
            EnsureSorted();
            uint32_t pos = location.start_offset;

            auto it = std::lower_bound(
                entries_.begin(), entries_.end(), pos, [](const Entry& e, uint32_t p) { return e.start < p; });

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

            std::optional<SymbolId> result;
            uint32_t min_span = UINT32_MAX;

            for (; it != entries_.end() && it->start <= pos; ++it)
            {
                if (pos >= it->start && pos < it->end)
                {
                    uint32_t span = it->end - it->start;
                    if (span < min_span)
                    {
                        min_span = span;
                        result = it->symbol_id;
                    }
                }
            }

            return result;
        }

    private:
        struct Entry
        {
            uint32_t start;
            uint32_t end;
            SymbolId symbol_id;

            bool operator<(const Entry& other) const
            {
                if (start != other.start)
                    return start < other.start;
                return end > other.end;
            }
        };

        void EnsureSorted() const
        {
            if (needs_sort_)
            {
                std::sort(entries_.begin(), entries_.end());
                needs_sort_ = false;
            }
        }

        mutable std::vector<Entry> entries_;
        mutable bool needs_sort_ = false;
    };

} // namespace lsp::language::symbol