#include <XCEngine/UI/Core/UIContext.h>

#include <algorithm>

namespace XCEngine {
namespace UI {

namespace {

void MergeChange(
    std::vector<UIElementChange>& changes,
    UIElementId id,
    UIElementChangeKind kind,
    UIDirtyFlags dirtyFlags) {
    for (UIElementChange& change : changes) {
        if (change.id != id) {
            continue;
        }

        if (change.kind == UIElementChangeKind::Created || kind == UIElementChangeKind::Created) {
            change.kind = UIElementChangeKind::Created;
        } else if (change.kind == UIElementChangeKind::Removed || kind == UIElementChangeKind::Removed) {
            change.kind = UIElementChangeKind::Removed;
        } else {
            change.kind = UIElementChangeKind::Updated;
        }
        change.dirtyFlags |= dirtyFlags;
        return;
    }

    UIElementChange change = {};
    change.id = id;
    change.kind = kind;
    change.dirtyFlags = dirtyFlags;
    changes.push_back(change);
}

bool IsStructureAffectingChange(UIDirtyFlags flags) {
    return HasAnyDirtyFlags(flags, UIDirtyFlags::Structure | UIDirtyFlags::Layout);
}

} // namespace

const UIElementChange* UIElementTreeRebuildResult::FindChange(UIElementId id) const {
    for (const UIElementChange& change : changes) {
        if (change.id == id) {
            return &change;
        }
    }

    return nullptr;
}

bool UIElementTreeRebuildResult::HasChange(UIElementId id) const {
    return FindChange(id) != nullptr;
}

UIElementTreeRebuildResult UIElementTree::Rebuild(const UIBuildList& buildList) {
    UIElementTreeRebuildResult result = {};
    result.generation = m_generation + 1;

    if (buildList.Empty()) {
        if (!m_nodes.empty()) {
            result.treeChanged = true;
            for (const auto& entry : m_nodes) {
                MergeChange(
                    result.changes,
                    entry.first,
                    UIElementChangeKind::Removed,
                    UIDirtyFlags::Structure | UIDirtyFlags::Layout | UIDirtyFlags::Paint);
            }
        }

        m_nodes.clear();
        m_rootId = kInvalidUIElementId;
        m_generation = result.generation;
        return result;
    }

    std::unordered_map<UIElementId, UIBuildElement> newElementsById;
    newElementsById.reserve(buildList.GetCount());
    for (const UIBuildElement& element : buildList.GetElements()) {
        if (element.id == kInvalidUIElementId) {
            result.succeeded = false;
            result.errorMessage = "UIElementTree::Rebuild received an invalid element id.";
            return result;
        }

        const auto insertResult = newElementsById.emplace(element.id, element);
        if (!insertResult.second) {
            result.succeeded = false;
            result.errorMessage = "UIElementTree::Rebuild received duplicate element ids.";
            return result;
        }
    }

    const UIElementId newRootId = buildList.GetElements().front().id;
    if (buildList.GetElements().front().parentId != kInvalidUIElementId) {
        result.succeeded = false;
        result.errorMessage = "UIElementTree::Rebuild requires the first element to be the root.";
        return result;
    }

    std::unordered_map<UIElementId, UIElementNode> newNodes;
    newNodes.reserve(buildList.GetCount());
    for (const UIBuildElement& element : buildList.GetElements()) {
        if (element.parentId != kInvalidUIElementId &&
            newElementsById.find(element.parentId) == newElementsById.end()) {
            result.succeeded = false;
            result.errorMessage = "UIElementTree::Rebuild found an element whose parent does not exist.";
            return result;
        }

        UIElementNode node = {};
        node.id = element.id;
        node.parentId = element.parentId;
        node.typeName = element.typeName;
        node.structuralRevision = element.structuralRevision;
        node.localStateRevision = element.localStateRevision;
        node.viewModelRevision = element.viewModelRevision;
        node.lastBuildGeneration = result.generation;

        const auto existingNode = m_nodes.find(element.id);
        if (existingNode != m_nodes.end()) {
            node.dirtyFlags = existingNode->second.dirtyFlags;
        }

        if (node.parentId != kInvalidUIElementId) {
            const auto parentNode = newNodes.find(node.parentId);
            if (parentNode == newNodes.end()) {
                result.succeeded = false;
                result.errorMessage = "UIElementTree::Rebuild requires parents to appear before children.";
                return result;
            }

            node.depth = parentNode->second.depth + 1u;
            parentNode->second.childIds.push_back(node.id);
        }

        newNodes.emplace(node.id, std::move(node));
    }

    for (const auto& oldEntry : m_nodes) {
        if (newNodes.find(oldEntry.first) != newNodes.end()) {
            continue;
        }

        result.treeChanged = true;
        MergeChange(
            result.changes,
            oldEntry.first,
            UIElementChangeKind::Removed,
            UIDirtyFlags::Structure | UIDirtyFlags::Layout | UIDirtyFlags::Paint);

        const UIElementId parentId = oldEntry.second.parentId;
        if (parentId != kInvalidUIElementId && newNodes.find(parentId) != newNodes.end()) {
            MarkDirtyInternal(newNodes, parentId, UIDirtyFlags::Structure | UIDirtyFlags::Layout | UIDirtyFlags::Paint);
        }
    }

    for (auto& newEntry : newNodes) {
        const UIElementId id = newEntry.first;
        UIElementNode& newNode = newEntry.second;

        const auto oldEntry = m_nodes.find(id);
        if (oldEntry == m_nodes.end()) {
            result.treeChanged = true;
            MarkDirtyInternal(newNodes, id, UIDirtyFlags::Structure | UIDirtyFlags::Layout | UIDirtyFlags::Paint);
            MergeChange(
                result.changes,
                id,
                UIElementChangeKind::Created,
                UIDirtyFlags::Structure | UIDirtyFlags::Layout | UIDirtyFlags::Paint);
            continue;
        }

        UIDirtyFlags dirtyFlags = UIDirtyFlags::None;
        const UIElementNode& oldNode = oldEntry->second;
        if (oldNode.parentId != newNode.parentId ||
            oldNode.typeName != newNode.typeName ||
            oldNode.structuralRevision != newNode.structuralRevision) {
            dirtyFlags |= UIDirtyFlags::Structure | UIDirtyFlags::Layout | UIDirtyFlags::Paint;
        }
        if (oldNode.localStateRevision != newNode.localStateRevision) {
            dirtyFlags |= UIDirtyFlags::LocalState | UIDirtyFlags::Paint;
        }
        if (oldNode.viewModelRevision != newNode.viewModelRevision) {
            dirtyFlags |= UIDirtyFlags::ViewModel | UIDirtyFlags::Paint;
        }

        if (oldNode.childIds != newNode.childIds) {
            dirtyFlags |= UIDirtyFlags::Structure | UIDirtyFlags::Layout | UIDirtyFlags::Paint;
        }

        if (!IsDirty(dirtyFlags)) {
            continue;
        }

        result.treeChanged = true;
        MarkDirtyInternal(newNodes, id, dirtyFlags);
        MergeChange(result.changes, id, UIElementChangeKind::Updated, dirtyFlags);
    }

    m_nodes = std::move(newNodes);
    m_rootId = newRootId;
    m_generation = result.generation;
    result.dirtyRootIds = CollectDirtyRootIds();
    return result;
}

bool UIElementTree::HasNode(UIElementId id) const {
    return m_nodes.find(id) != m_nodes.end();
}

const UIElementNode* UIElementTree::FindNode(UIElementId id) const {
    const auto entry = m_nodes.find(id);
    return entry != m_nodes.end()
        ? &entry->second
        : nullptr;
}

void UIElementTree::MarkDirty(UIElementId id, UIDirtyFlags flags) {
    MarkDirtyInternal(m_nodes, id, flags);
}

std::vector<UIElementId> UIElementTree::CollectDirtyRootIds() const {
    std::vector<UIElementId> dirtyRoots;
    dirtyRoots.reserve(m_nodes.size());

    for (const auto& entry : m_nodes) {
        const UIElementNode& node = entry.second;
        if (!node.IsDirty()) {
            continue;
        }

        const auto parentEntry = m_nodes.find(node.parentId);
        if (parentEntry != m_nodes.end() && parentEntry->second.IsDirty()) {
            continue;
        }

        dirtyRoots.push_back(node.id);
    }

    std::sort(
        dirtyRoots.begin(),
        dirtyRoots.end(),
        [this](UIElementId left, UIElementId right) {
            const UIElementNode* leftNode = FindNode(left);
            const UIElementNode* rightNode = FindNode(right);
            const std::uint32_t leftDepth = leftNode != nullptr ? leftNode->depth : 0u;
            const std::uint32_t rightDepth = rightNode != nullptr ? rightNode->depth : 0u;
            return leftDepth != rightDepth
                ? leftDepth < rightDepth
                : left < right;
        });
    return dirtyRoots;
}

void UIElementTree::ClearDirtyFlags(UIElementId id, bool includeSubtree) {
    if (!includeSubtree) {
        const auto entry = m_nodes.find(id);
        if (entry != m_nodes.end()) {
            entry->second.dirtyFlags = UIDirtyFlags::None;
        }
        return;
    }

    ClearDirtyFlagsRecursive(id);
}

void UIElementTree::ClearAllDirtyFlags() {
    for (auto& entry : m_nodes) {
        entry.second.dirtyFlags = UIDirtyFlags::None;
    }
}

void UIElementTree::MarkDirtyInternal(
    std::unordered_map<UIElementId, UIElementNode>& nodes,
    UIElementId id,
    UIDirtyFlags flags) {
    const auto entry = nodes.find(id);
    if (entry == nodes.end()) {
        return;
    }

    entry->second.dirtyFlags |= flags;
    if (!IsStructureAffectingChange(flags)) {
        return;
    }

    UIElementId parentId = entry->second.parentId;
    while (parentId != kInvalidUIElementId) {
        const auto parentEntry = nodes.find(parentId);
        if (parentEntry == nodes.end()) {
            break;
        }

        parentEntry->second.dirtyFlags |= UIDirtyFlags::Layout;
        parentId = parentEntry->second.parentId;
    }
}

void UIElementTree::ClearDirtyFlagsRecursive(UIElementId id) {
    const auto entry = m_nodes.find(id);
    if (entry == m_nodes.end()) {
        return;
    }

    entry->second.dirtyFlags = UIDirtyFlags::None;
    const std::vector<UIElementId> childIds = entry->second.childIds;
    for (UIElementId childId : childIds) {
        ClearDirtyFlagsRecursive(childId);
    }
}

UIElementTreeRebuildResult UIContext::Rebuild(const BuildCallback& buildCallback) {
    m_buildList.Reset();
    UIBuildContext buildContext(m_buildList);
    if (buildCallback) {
        buildCallback(buildContext);
    }

    UIElementTreeRebuildResult result = {};
    result.generation = m_tree.GetGeneration() + 1;

    if (!buildContext.IsValid()) {
        result.succeeded = false;
        result.errorMessage = buildContext.GetLastError();
        return result;
    }

    if (buildContext.HasOpenElements()) {
        result.succeeded = false;
        result.errorMessage = "UIContext::Rebuild finished with unclosed UI elements.";
        return result;
    }

    return m_tree.Rebuild(m_buildList);
}

} // namespace UI
} // namespace XCEngine