#include <unordered_set>
#include "./relations.hpp"

namespace lsp::language::symbol
{
    // ===== ReferenceGraph 实现 =====

    void ReferenceGraph::AddReference(SymbolId symbol_id, const ast::Location& location, bool is_definition, bool is_write)
    {
        references_[symbol_id].push_back({ location, symbol_id, is_definition, is_write });
    }

    const std::vector<Reference>* ReferenceGraph::GetReferences(SymbolId symbol_id) const
    {
        auto it = references_.find(symbol_id);
        return it != references_.end() ? &it->second : nullptr;
    }

    std::optional<ast::Location> ReferenceGraph::GetDefinition(SymbolId symbol_id) const
    {
        auto it = references_.find(symbol_id);
        if (it != references_.end())
        {
            for (const auto& ref : it->second)
            {
                if (ref.is_definition)
                    return ref.location;
            }
        }
        return std::nullopt;
    }

    void ReferenceGraph::Clear()
    {
        references_.clear();
    }

    // ===== InheritanceGraph 实现 =====

    void InheritanceGraph::AddInheritance(SymbolId derived_class, SymbolId base_class)
    {
        base_classes_[derived_class].push_back(base_class);
        derived_classes_[base_class].push_back(derived_class);
    }

    const std::vector<SymbolId>* InheritanceGraph::GetBaseClasses(SymbolId class_id) const
    {
        auto it = base_classes_.find(class_id);
        return it != base_classes_.end() ? &it->second : nullptr;
    }

    const std::vector<SymbolId>* InheritanceGraph::GetDerivedClasses(SymbolId class_id) const
    {
        auto it = derived_classes_.find(class_id);
        return it != derived_classes_.end() ? &it->second : nullptr;
    }

    bool InheritanceGraph::IsSubclassOf(SymbolId derived, SymbolId base) const
    {
        std::unordered_set<SymbolId> visited;
        std::vector<SymbolId> queue = { derived };

        while (!queue.empty())
        {
            SymbolId current = queue.back();
            queue.pop_back();

            if (current == base)
                return true;

            if (visited.count(current))
                continue;
            visited.insert(current);

            auto bases = GetBaseClasses(current);
            if (bases)
                queue.insert(queue.end(), bases->begin(), bases->end());
        }

        return false;
    }

    void InheritanceGraph::Clear()
    {
        base_classes_.clear();
        derived_classes_.clear();
    }

    // ===== CallGraph 实现 =====

    void CallGraph::AddCall(SymbolId caller, SymbolId callee, const ast::Location& location)
    {
        CallRelation rel{ caller, callee, location };
        callers_[callee].push_back(rel);
        callees_[caller].push_back(rel);
    }

    const std::vector<CallRelation>* CallGraph::GetCallers(SymbolId function_id) const
    {
        auto it = callers_.find(function_id);
        return it != callers_.end() ? &it->second : nullptr;
    }

    const std::vector<CallRelation>* CallGraph::GetCallees(SymbolId function_id) const
    {
        auto it = callees_.find(function_id);
        return it != callees_.end() ? &it->second : nullptr;
    }

    void CallGraph::Clear()
    {
        callers_.clear();
        callees_.clear();
    }

}
