#include "TableTreeBuilder.h"
#include "catalog/PgInheritsContainer.h"
#include "TableNode.h"
#include <ranges>

TableTreeBuilder::TableTreeBuilder(
    const std::map<Oid, PgClass> &source,
    const IFindParents &inheritance
)
: source(source)
, inheritance(inheritance)
{
}

std::shared_ptr<TableNode> TableTreeBuilder::Build()
{
    rootNode = std::make_shared<TableNode>();

    // when childrens are ordered before there parents
    // the parent will automatically be added first (recursively)
    // processed nodes are tracked in the nodes member
    for (auto && e : source | std::views::filter(
        [this] (auto &e)
        {
            return processedNodes.count(e.first) == 0;
        })
    )
    {
        addNode(e.second);
    }
    AssignNodeIndexesAndParents(rootNode);
    auto ret = rootNode;
    resetState();
    return ret;
}

void TableTreeBuilder::resetState()
{
    rootNode.reset();
    processedNodes.clear();
}

std::shared_ptr<TableNode>  TableTreeBuilder::addNode(const PgClass &cls)
{
    std::shared_ptr<TableNode> node = std::make_shared<TableNode>(cls);
    processedNodes.emplace(cls.oid(), node);
    std::vector<Oid> parents = inheritance.getParentsOf(cls.oid());
    if (parents.empty())
        addToplevelNode(node);
    else
        addToParents(node, parents);
    return node;
}

void TableTreeBuilder::addToplevelNode(std::shared_ptr<TableNode> node)
{
    rootNode->children.push_back(node);
}

void TableTreeBuilder::addToParents(std::shared_ptr<TableNode> node, const std::vector<Oid> &parents)
{
    for (auto &parent : parents)
    {
        getParent(parent)->children.push_back(node);
    }
}

std::shared_ptr<TableNode> TableTreeBuilder::getParent(Oid oid)
{
    auto parent = processedNodes.find(oid);
    if (parent != processedNodes.end())
        return parent->second;

    // Not present, find in source list and add now
    auto source_iter = source.find(oid);
    assert(source_iter != source.end());
    return addNode(source_iter->second);
}

void TableTreeBuilder::AssignNodeIndexesAndParents(std::shared_ptr<TableNode> parent)
{
    int index = 0;
    for (auto &n : parent->children)
    {
        n->parent = parent;
        n->myIndex = index++;
        AssignNodeIndexesAndParents(n);
    }
}