#include "treebuilder.h"
#include "xtreewidget.h"
#include "xtreewidgetitem.h"
#include "typeregistry.h"
#include "logmanager.h"

#include <QCollator>
#include <algorithm>

// Natural sort comparator for strings with numbers (e.g., "chunk1" < "chunk2" < "chunk10")
static bool naturalLessThan(const QString& a, const QString& b) {
    static QCollator collator;
    collator.setNumericMode(true);
    collator.setCaseSensitivity(Qt::CaseInsensitive);
    return collator.compare(a, b) < 0;
}

TreeBuilder::TreeBuilder(XTreeWidget* tree, const TypeRegistry& registry)
    : m_tree(tree)
    , m_registry(&registry)
{
}

void TreeBuilder::reset()
{
    m_categoryRoots.clear();
}

QString TreeBuilder::pluralizeType(const QString& typeName) const
{
    const auto& mod = m_registry->module();
    const auto it = mod.types.find(typeName);
    QString groupLabel = (it != mod.types.end() && !it->display.isEmpty())
                             ? it->display
                             : typeName;
    // Don't double the 's' if already ends with 's'
    if (groupLabel.endsWith('s') || groupLabel.endsWith('S'))
        return groupLabel;
    return groupLabel + "s";
}

XTreeWidgetItem* TreeBuilder::ensureTypeCategoryRoot(const QString& typeName, const QString& displayOverride)
{
    const QString categoryKey = displayOverride.isEmpty() ? typeName : displayOverride;

    if (m_categoryRoots.contains(categoryKey))
        return m_categoryRoots[categoryKey];

    auto* root = new XTreeWidgetItem(m_tree);
    QString categoryLabel;
    if (displayOverride.isEmpty()) {
        categoryLabel = pluralizeType(typeName);
    } else if (displayOverride.endsWith('s') || displayOverride.endsWith('S')) {
        categoryLabel = displayOverride;  // Don't double the 's'
    } else {
        categoryLabel = displayOverride + "s";
    }
    root->setText(0, categoryLabel);
    root->setData(0, Qt::UserRole + 1, "CATEGORY");
    root->setData(0, Qt::UserRole + 2, typeName);
    m_tree->addTopLevelItem(root);

    m_categoryRoots.insert(categoryKey, root);
    return root;
}

XTreeWidgetItem* TreeBuilder::ensureSubcategory(XTreeWidgetItem* parent, const QString& childTypeName)
{
    // Look for existing subcategory
    for (int i = 0; i < parent->childCount(); i++) {
        auto* c = static_cast<XTreeWidgetItem*>(parent->child(i));
        if (c->data(0, Qt::UserRole + 1).toString() == "SUBCATEGORY" &&
            c->data(0, Qt::UserRole + 2).toString() == childTypeName)
        {
            return c;
        }
    }

    auto* sub = new XTreeWidgetItem(parent);
    sub->setText(0, pluralizeType(childTypeName));
    sub->setData(0, Qt::UserRole + 1, "SUBCATEGORY");
    sub->setData(0, Qt::UserRole + 2, childTypeName);
    parent->addChild(sub);
    sub->setExpanded(false);
    return sub;
}

XTreeWidgetItem* TreeBuilder::addInstanceNode(XTreeWidgetItem* parent, const QString& displayName,
                                               const QString& typeName, const QVariantMap& vars)
{
    auto* inst = new XTreeWidgetItem(parent);
    inst->setText(0, displayName);
    inst->setData(0, Qt::UserRole + 1, "INSTANCE");
    inst->setData(0, Qt::UserRole + 2, typeName);
    inst->setData(0, Qt::UserRole + 3, vars);
    parent->addChild(inst);
    inst->setExpanded(false);
    return inst;
}

QString TreeBuilder::instanceDisplayFor(const QVariantMap& obj, const QString& fallbackType,
                                         const QString& fallbackKey, std::optional<int> index)
{
    // _name takes priority (explicit name set by script via set_name())
    if (DslKeys::contains(obj, DslKey::Name)) {
        const QString s = DslKeys::getString(obj, DslKey::Name);
        if (!s.isEmpty()) return s;
    }

    // _display is secondary (set via set_display())
    if (DslKeys::contains(obj, DslKey::Display)) {
        const QString s = DslKeys::getString(obj, DslKey::Display);
        if (!s.isEmpty()) return s;
    }

    // Index-based fallback for array items - ensures unique names
    // This comes BEFORE type fallback to avoid all items showing same type name
    if (!fallbackKey.isEmpty() && index.has_value()) {
        return QString("%1[%2]").arg(fallbackKey).arg(*index);
    }

    // Type fallback (for non-array items)
    if (!fallbackType.isEmpty()) return fallbackType;

    if (!fallbackKey.isEmpty()) {
        return fallbackKey;
    }
    return QStringLiteral("<unnamed>");
}

void TreeBuilder::addParsedFile(const QString& typeName, const QVariantMap& vars, const QString& fileName)
{
    XTreeWidgetItem* cat = ensureTypeCategoryRoot(typeName, DslKeys::getString(vars, DslKey::Display));
    const QString displayName = DslKeys::contains(vars, DslKey::Name)
        ? DslKeys::getString(vars, DslKey::Name)
        : fileName;
    XTreeWidgetItem* inst = addInstanceNode(cat, displayName, typeName, vars);
    routeNestedObjects(inst, vars);
}

void TreeBuilder::routeNestedObjects(XTreeWidgetItem* parent, const QVariantMap& vars)
{
    for (auto it = vars.begin(); it != vars.end(); ++it) {
        const QString& key = it.key();
        const QVariant& v = it.value();

        // Child object (QVariantMap with _type)
        if (v.typeId() == QMetaType::QVariantMap) {
            const QVariantMap child = v.toMap();

            // Skip hidden objects
            if (DslKeys::get(child, DslKey::Hidden).toBool()) {
                continue;
            }

            const QString childType = DslKeys::getString(child, DslKey::Type);
            if (!childType.isEmpty()) {
                auto* subcat = ensureSubcategory(parent, childType);
                const QString childDisplay = instanceDisplayFor(child, childType, key);
                auto* childInst = addInstanceNode(subcat, childDisplay, childType, child);
                routeNestedObjects(childInst, child);
            }
            continue;
        }

        // Array of child objects
        if (v.typeId() == QMetaType::QVariantList) {
            // Check for skip marker
            if (DslKeys::hasSkipTree(vars, key)) {
                LogManager::instance().debug(QString("[TREE] Skipping array '%1' (has skip_tree marker)").arg(key));
                continue;
            }

            const QVariantList list = v.toList();
            LogManager::instance().debug(QString("[TREE] Processing array '%1' with %2 items")
                .arg(key).arg(list.size()));

            for (int i = 0; i < list.size(); i++) {
                if (list[i].typeId() != QMetaType::QVariantMap) {
                    continue;
                }
                const QVariantMap child = list[i].toMap();

                // Skip hidden objects
                if (DslKeys::get(child, DslKey::Hidden).toBool()) {
                    continue;
                }

                const QString childType = DslKeys::getString(child, DslKey::Type);
                if (childType.isEmpty()) {
                    continue;
                }

                auto* subcat = ensureSubcategory(parent, childType);
                const QString childDisplay = instanceDisplayFor(child, childType, key, i);
                auto* childInst = addInstanceNode(subcat, childDisplay, childType, child);
                routeNestedObjects(childInst, child);
            }
        }
    }
}

void TreeBuilder::organizeChildrenByExtension(XTreeWidgetItem* parent)
{
    if (!parent) return;

    // Recursively process children first (depth-first)
    for (int i = 0; i < parent->childCount(); i++) {
        organizeChildrenByExtension(static_cast<XTreeWidgetItem*>(parent->child(i)));
    }

    // Only organize SUBCATEGORY nodes
    QString nodeType = parent->data(0, Qt::UserRole + 1).toString();
    if (nodeType != "SUBCATEGORY") return;

    // Group children by extension
    QMap<QString, QList<XTreeWidgetItem*>> byExtension;
    QList<XTreeWidgetItem*> noExtension;
    QList<XTreeWidgetItem*> nonInstanceChildren;

    // Take all children
    QList<XTreeWidgetItem*> children;
    while (parent->childCount() > 0) {
        children.append(static_cast<XTreeWidgetItem*>(parent->takeChild(0)));
    }

    for (auto* child : children) {
        QString childNodeType = child->data(0, Qt::UserRole + 1).toString();

        // Keep non-instance children as-is
        if (childNodeType == "SUBCATEGORY" || childNodeType == "EXTENSION_GROUP") {
            nonInstanceChildren.append(child);
            continue;
        }

        QVariantMap vars = child->data(0, Qt::UserRole + 3).toMap();
        QString name = DslKeys::getString(vars, DslKey::Name);

        // Skip names starting with dot (indexed chunk names)
        if (name.startsWith('.')) {
            noExtension.append(child);
            continue;
        }

        int dotPos = name.lastIndexOf('.');
        if (dotPos > 0) {
            QString ext = name.mid(dotPos + 1).toLower();
            byExtension[ext].append(child);
        } else {
            noExtension.append(child);
        }
    }

    // Re-add non-instance children first
    for (auto* child : nonInstanceChildren) {
        parent->addChild(child);
    }

    // Only create groups if multiple extensions exist
    int uniqueExtensions = byExtension.size() + (noExtension.isEmpty() ? 0 : 1);
    if (uniqueExtensions <= 1) {
        // Put all items back directly (sorted)
        QList<XTreeWidgetItem*> allItems;
        for (auto& list : byExtension) {
            allItems.append(list);
        }
        allItems.append(noExtension);

        std::sort(allItems.begin(), allItems.end(), [](XTreeWidgetItem* a, XTreeWidgetItem* b) {
            return naturalLessThan(a->text(0), b->text(0));
        });

        for (auto* child : allItems) {
            parent->addChild(child);
        }
        return;
    }

    // Create extension groups
    QStringList sortedExts = byExtension.keys();
    std::sort(sortedExts.begin(), sortedExts.end());

    for (const QString& ext : sortedExts) {
        auto* extGroup = new XTreeWidgetItem(parent);
        extGroup->setText(0, QString(".%1").arg(ext));
        extGroup->setData(0, Qt::UserRole + 1, "EXTENSION_GROUP");
        extGroup->setData(0, Qt::UserRole + 2, ext);

        QList<XTreeWidgetItem*>& items = byExtension[ext];
        std::sort(items.begin(), items.end(), [](XTreeWidgetItem* a, XTreeWidgetItem* b) {
            return naturalLessThan(a->text(0), b->text(0));
        });

        for (auto* item : items) {
            extGroup->addChild(item);
        }
        extGroup->setExpanded(false);
    }

    // Add items with no extension at the end
    if (!noExtension.isEmpty()) {
        auto* otherGroup = new XTreeWidgetItem(parent);
        otherGroup->setText(0, "(other)");
        otherGroup->setData(0, Qt::UserRole + 1, "EXTENSION_GROUP");
        otherGroup->setData(0, Qt::UserRole + 2, "");

        std::sort(noExtension.begin(), noExtension.end(), [](XTreeWidgetItem* a, XTreeWidgetItem* b) {
            return naturalLessThan(a->text(0), b->text(0));
        });

        for (auto* item : noExtension) {
            otherGroup->addChild(item);
        }
        otherGroup->setExpanded(false);
    }
}

void TreeBuilder::updateNodeCounts(XTreeWidgetItem* node)
{
    if (!node) return;

    // Recursively update children first
    for (int i = 0; i < node->childCount(); i++) {
        updateNodeCounts(static_cast<XTreeWidgetItem*>(node->child(i)));
    }

    // Update count for grouping nodes
    QString nodeType = node->data(0, Qt::UserRole + 1).toString();
    if (nodeType == "SUBCATEGORY" || nodeType == "CATEGORY" || nodeType == "EXTENSION_GROUP") {
        int count = node->childCount();
        if (count > 0) {
            QString currentText = node->text(0);
            int parenPos = currentText.lastIndexOf(" (");
            if (parenPos > 0) {
                currentText = currentText.left(parenPos);
            }
            node->setText(0, QString("%1 (%2)").arg(currentText).arg(count));
        }
    }
}