#include "ui_core/node_lifetime.h"

#include <algorithm>
#include <utility>

namespace pp::ui {
namespace {

void bump_generation(std::uint32_t& generation) noexcept
{
    ++generation;
    if (generation == 0U) {
        generation = 1U;
    }
}

}

ScopedUiConnection::ScopedUiConnection(NodeLifetimeTree& tree, UiConnection connection) noexcept
    : tree_(&tree)
    , connection_(connection)
{
}

ScopedUiConnection::ScopedUiConnection(ScopedUiConnection&& other) noexcept
    : tree_(std::exchange(other.tree_, nullptr))
    , connection_(std::exchange(other.connection_, UiConnection {}))
{
}

ScopedUiConnection& ScopedUiConnection::operator=(ScopedUiConnection&& other) noexcept
{
    if (this != &other) {
        reset();
        tree_ = std::exchange(other.tree_, nullptr);
        connection_ = std::exchange(other.connection_, UiConnection {});
    }

    return *this;
}

ScopedUiConnection::~ScopedUiConnection()
{
    reset();
}

void ScopedUiConnection::reset() noexcept
{
    if (tree_ != nullptr && connection_.valid()) {
        (void)tree_->disconnect(connection_);
    }
    tree_ = nullptr;
    connection_ = {};
}

UiConnection ScopedUiConnection::connection() const noexcept
{
    return connection_;
}

bool ScopedUiConnection::connected() const noexcept
{
    return tree_ != nullptr && connection_.valid();
}

pp::foundation::Result<NodeHandle> NodeLifetimeTree::create_root()
{
    return allocate_node(NodeHandle {});
}

pp::foundation::Result<NodeHandle> NodeLifetimeTree::create_child(NodeHandle parent)
{
    if (node_slot(parent) == nullptr) {
        return pp::foundation::Result<NodeHandle>::failure(
            pp::foundation::Status::invalid_argument("UI child node requires a live parent"));
    }

    auto child = allocate_node(parent);
    if (!child) {
        return child;
    }

    node_slot(parent)->children.push_back(child.value());
    return child;
}

bool NodeLifetimeTree::contains(NodeHandle node) const noexcept
{
    return node_slot(node) != nullptr;
}

pp::foundation::Result<NodeHandle> NodeLifetimeTree::parent_of(NodeHandle node) const noexcept
{
    const auto* slot = node_slot(node);
    if (slot == nullptr) {
        return pp::foundation::Result<NodeHandle>::failure(
            pp::foundation::Status::invalid_argument("UI node handle is not live"));
    }

    if (!slot->parent.valid()) {
        return pp::foundation::Result<NodeHandle>::failure(
            pp::foundation::Status::invalid_argument("UI root node does not have a parent"));
    }

    return pp::foundation::Result<NodeHandle>::success(slot->parent);
}

pp::foundation::Result<std::size_t> NodeLifetimeTree::child_count(NodeHandle node) const noexcept
{
    const auto* slot = node_slot(node);
    if (slot == nullptr) {
        return pp::foundation::Result<std::size_t>::failure(
            pp::foundation::Status::invalid_argument("UI node handle is not live"));
    }

    return pp::foundation::Result<std::size_t>::success(slot->children.size());
}

pp::foundation::Result<NodeHandle> NodeLifetimeTree::child_at(NodeHandle node, std::size_t index) const noexcept
{
    const auto* slot = node_slot(node);
    if (slot == nullptr) {
        return pp::foundation::Result<NodeHandle>::failure(
            pp::foundation::Status::invalid_argument("UI node handle is not live"));
    }

    if (index >= slot->children.size()) {
        return pp::foundation::Result<NodeHandle>::failure(
            pp::foundation::Status::out_of_range("UI child index is outside the node"));
    }

    return pp::foundation::Result<NodeHandle>::success(slot->children[index]);
}

pp::foundation::Status NodeLifetimeTree::destroy_subtree(NodeHandle node) noexcept
{
    auto* slot = node_slot(node);
    if (slot == nullptr) {
        return pp::foundation::Status::invalid_argument("UI node handle is not live");
    }

    const auto children = slot->children;
    for (const auto child : children) {
        (void)destroy_subtree(child);
    }

    const auto connections = slot->connections;
    for (const auto connection : connections) {
        release_connection(connection);
    }

    unlink_from_parent(node);
    slot = node_slot(node);
    if (slot == nullptr) {
        return pp::foundation::Status::success();
    }

    slot->children.clear();
    slot->connections.clear();
    slot->parent = {};
    release_captures_for_node(node);
    slot->alive = false;
    free_nodes_.push_back(node.slot);
    return pp::foundation::Status::success();
}

void NodeLifetimeTree::clear() noexcept
{
    free_nodes_.clear();
    for (std::uint32_t index = 0; index < nodes_.size(); ++index) {
        auto& slot = nodes_[index];
        slot.alive = false;
        slot.parent = {};
        slot.children.clear();
        slot.connections.clear();
        free_nodes_.push_back(index);
    }

    free_connections_.clear();
    for (std::uint32_t index = 0; index < connections_.size(); ++index) {
        auto& slot = connections_[index];
        slot.alive = false;
        slot.node = {};
        slot.callback = nullptr;
        free_connections_.push_back(index);
    }

    captures_ = {};
}

pp::foundation::Status NodeLifetimeTree::capture(UiCaptureKind kind, NodeHandle node) noexcept
{
    if (node_slot(node) == nullptr) {
        return pp::foundation::Status::invalid_argument("UI capture requires a live node");
    }

    captures_[capture_index(kind)] = node;
    return pp::foundation::Status::success();
}

pp::foundation::Status NodeLifetimeTree::release_capture(UiCaptureKind kind, NodeHandle node) noexcept
{
    auto& captured = captures_[capture_index(kind)];
    if (captured != node) {
        return pp::foundation::Status::invalid_argument("UI capture release requires the captured node");
    }

    captured = {};
    return pp::foundation::Status::success();
}

pp::foundation::Result<NodeHandle> NodeLifetimeTree::captured_node(UiCaptureKind kind) const noexcept
{
    const auto captured = captures_[capture_index(kind)];
    if (node_slot(captured) == nullptr) {
        return pp::foundation::Result<NodeHandle>::failure(
            pp::foundation::Status::invalid_argument("UI capture slot is empty"));
    }

    return pp::foundation::Result<NodeHandle>::success(captured);
}

pp::foundation::Result<UiConnection> NodeLifetimeTree::connect(NodeHandle node, Callback callback)
{
    if (node_slot(node) == nullptr) {
        return pp::foundation::Result<UiConnection>::failure(
            pp::foundation::Status::invalid_argument("UI connection requires a live node"));
    }

    if (!callback) {
        return pp::foundation::Result<UiConnection>::failure(
            pp::foundation::Status::invalid_argument("UI connection requires a callback"));
    }

    UiConnection connection;
    if (!free_connections_.empty()) {
        connection.slot = free_connections_.back();
        free_connections_.pop_back();
        auto& slot = connections_[connection.slot];
        bump_generation(slot.generation);
        slot.alive = true;
        slot.node = node;
        slot.callback = std::move(callback);
        connection.generation = slot.generation;
    } else {
        connection.slot = static_cast<std::uint32_t>(connections_.size());
        connection.generation = 1U;
        connections_.push_back(ConnectionSlot {
            .generation = connection.generation,
            .alive = true,
            .node = node,
            .callback = std::move(callback),
        });
    }

    node_slot(node)->connections.push_back(connection);
    return pp::foundation::Result<UiConnection>::success(connection);
}

pp::foundation::Result<ScopedUiConnection> NodeLifetimeTree::scoped_connect(NodeHandle node, Callback callback)
{
    auto connection = connect(node, std::move(callback));
    if (!connection) {
        return pp::foundation::Result<ScopedUiConnection>::failure(connection.status());
    }

    return pp::foundation::Result<ScopedUiConnection>::success(ScopedUiConnection(*this, connection.value()));
}

pp::foundation::Status NodeLifetimeTree::disconnect(UiConnection connection) noexcept
{
    if (connection_slot(connection) == nullptr) {
        return pp::foundation::Status::invalid_argument("UI connection handle is not live");
    }

    release_connection(connection);
    return pp::foundation::Status::success();
}

pp::foundation::Result<std::size_t> NodeLifetimeTree::live_connection_count(NodeHandle node) const noexcept
{
    const auto* slot = node_slot(node);
    if (slot == nullptr) {
        return pp::foundation::Result<std::size_t>::failure(
            pp::foundation::Status::invalid_argument("UI node handle is not live"));
    }

    return pp::foundation::Result<std::size_t>::success(slot->connections.size());
}

pp::foundation::Status NodeLifetimeTree::dispatch(NodeHandle node)
{
    const auto* slot = node_slot(node);
    if (slot == nullptr) {
        return pp::foundation::Status::invalid_argument("UI dispatch requires a live node");
    }

    const auto connections = slot->connections;
    for (const auto connection : connections) {
        if (node_slot(node) == nullptr) {
            break;
        }

        auto* connection_value = connection_slot(connection);
        if (connection_value == nullptr || connection_value->node != node) {
            continue;
        }

        connection_value->callback(node);
    }

    return pp::foundation::Status::success();
}

NodeLifetimeTree::NodeSlot* NodeLifetimeTree::node_slot(NodeHandle node) noexcept
{
    if (!node.valid() || node.slot >= nodes_.size()) {
        return nullptr;
    }

    auto& slot = nodes_[node.slot];
    if (!slot.alive || slot.generation != node.generation) {
        return nullptr;
    }

    return &slot;
}

const NodeLifetimeTree::NodeSlot* NodeLifetimeTree::node_slot(NodeHandle node) const noexcept
{
    if (!node.valid() || node.slot >= nodes_.size()) {
        return nullptr;
    }

    const auto& slot = nodes_[node.slot];
    if (!slot.alive || slot.generation != node.generation) {
        return nullptr;
    }

    return &slot;
}

NodeLifetimeTree::ConnectionSlot* NodeLifetimeTree::connection_slot(UiConnection connection) noexcept
{
    if (!connection.valid() || connection.slot >= connections_.size()) {
        return nullptr;
    }

    auto& slot = connections_[connection.slot];
    if (!slot.alive || slot.generation != connection.generation) {
        return nullptr;
    }

    return &slot;
}

const NodeLifetimeTree::ConnectionSlot* NodeLifetimeTree::connection_slot(UiConnection connection) const noexcept
{
    if (!connection.valid() || connection.slot >= connections_.size()) {
        return nullptr;
    }

    const auto& slot = connections_[connection.slot];
    if (!slot.alive || slot.generation != connection.generation) {
        return nullptr;
    }

    return &slot;
}

pp::foundation::Result<NodeHandle> NodeLifetimeTree::allocate_node(NodeHandle parent)
{
    NodeHandle node;
    if (!free_nodes_.empty()) {
        node.slot = free_nodes_.back();
        free_nodes_.pop_back();
        auto& slot = nodes_[node.slot];
        bump_generation(slot.generation);
        slot.alive = true;
        slot.parent = parent;
        slot.children.clear();
        slot.connections.clear();
        node.generation = slot.generation;
    } else {
        node.slot = static_cast<std::uint32_t>(nodes_.size());
        node.generation = 1U;
        nodes_.push_back(NodeSlot {
            .generation = node.generation,
            .alive = true,
            .parent = parent,
            .children = {},
            .connections = {},
        });
    }

    return pp::foundation::Result<NodeHandle>::success(node);
}

std::size_t NodeLifetimeTree::capture_index(UiCaptureKind kind) const noexcept
{
    switch (kind) {
    case UiCaptureKind::pointer:
        return 0U;
    case UiCaptureKind::keyboard:
        return 1U;
    }

    return 0U;
}

void NodeLifetimeTree::unlink_from_parent(NodeHandle node) noexcept
{
    const auto* slot = node_slot(node);
    if (slot == nullptr || !slot->parent.valid()) {
        return;
    }

    auto* parent = node_slot(slot->parent);
    if (parent == nullptr) {
        return;
    }

    parent->children.erase(
        std::remove(parent->children.begin(), parent->children.end(), node),
        parent->children.end());
}

void NodeLifetimeTree::release_connection(UiConnection connection) noexcept
{
    auto* slot = connection_slot(connection);
    if (slot == nullptr) {
        return;
    }

    auto* node = node_slot(slot->node);
    if (node != nullptr) {
        node->connections.erase(
            std::remove(node->connections.begin(), node->connections.end(), connection),
            node->connections.end());
    }

    slot->alive = false;
    slot->node = {};
    slot->callback = nullptr;
    free_connections_.push_back(connection.slot);
}

void NodeLifetimeTree::release_captures_for_node(NodeHandle node) noexcept
{
    for (auto& capture : captures_) {
        if (capture == node) {
            capture = {};
        }
    }
}

}