#include "ui/FolderTreeModel.h"
#include <QDir>

namespace whittle {

FolderTreeModel::FolderTreeModel(QObject* parent) : QFileSystemModel(parent) {
    setRootPath(QDir::rootPath());
    setFilter(QDir::AllDirs | QDir::NoDotAndDotDot | QDir::Hidden);
}

bool FolderTreeModel::effectivelyChecked(const QString& path) const {
    if (!m_recursive) return m_checked.contains(path);

    int  bestLen = -1;
    bool checked = false;
    const auto consider = [&](const QString& p, bool isChecked) {
        if (p == path || path.startsWith(p + '/')) {
            if (p.length() > bestLen) { bestLen = p.length(); checked = isChecked; }
        }
    };
    for (const QString& r : m_checked)  consider(r, true);
    for (const QString& e : m_excluded) consider(e, false);
    return bestLen >= 0 && checked;
}

Qt::CheckState FolderTreeModel::stateFor(const QString& path) const {
    const QString prefix = path + '/';
    if (effectivelyChecked(path)) {
        for (const QString& e : m_excluded)
            if (e.startsWith(prefix)) return Qt::PartiallyChecked;
        return Qt::Checked;
    }
    for (const QString& r : m_checked)
        if (r.startsWith(prefix)) return Qt::PartiallyChecked;
    return Qt::Unchecked;
}

void FolderTreeModel::check(const QString& path) {
    const QString prefix = path + '/';
    m_excluded.remove(path);
    for (auto it = m_excluded.begin(); it != m_excluded.end(); )
        it = it->startsWith(prefix) ? m_excluded.erase(it) : ++it;

    if (m_recursive) {
        if (effectivelyChecked(path)) return;
        for (auto it = m_checked.begin(); it != m_checked.end(); )
            it = it->startsWith(prefix) ? m_checked.erase(it) : ++it;
    }
    m_checked.insert(path);
}

void FolderTreeModel::uncheck(const QString& path) {
    const QString prefix = path + '/';
    m_checked.remove(path);
    if (!m_recursive) return;

    for (auto it = m_checked.begin(); it != m_checked.end(); )
        it = it->startsWith(prefix) ? m_checked.erase(it) : ++it;
    for (auto it = m_excluded.begin(); it != m_excluded.end(); )
        it = it->startsWith(prefix) ? m_excluded.erase(it) : ++it;

    if (effectivelyChecked(path)) m_excluded.insert(path);
}

Qt::ItemFlags FolderTreeModel::flags(const QModelIndex& ix) const {
    Qt::ItemFlags f = QFileSystemModel::flags(ix);
    if (ix.column() == 0 && isDir(ix))
        f |= Qt::ItemIsUserCheckable;
    return f;
}

QVariant FolderTreeModel::data(const QModelIndex& ix, int role) const {
    if (role == Qt::CheckStateRole && ix.column() == 0 && isDir(ix))
        return stateFor(filePath(ix));
    return QFileSystemModel::data(ix, role);
}

bool FolderTreeModel::setData(const QModelIndex& ix, const QVariant& v, int role) {
    if (role == Qt::CheckStateRole && ix.column() == 0 && isDir(ix)) {
        if (v.toInt() == Qt::Checked) check(filePath(ix));
        else                          uncheck(filePath(ix));
        emitSubtree(ix);
        emitAncestors(ix);
        emit selectionChanged();
        return true;
    }
    return QFileSystemModel::setData(ix, v, role);
}

void FolderTreeModel::emitSubtree(const QModelIndex& ix) {
    if (ix.isValid())
        emit dataChanged(ix, ix, {Qt::CheckStateRole});
    const int rows = rowCount(ix);
    for (int r = 0; r < rows; ++r)
        emitSubtree(index(r, 0, ix));
}

void FolderTreeModel::emitAncestors(const QModelIndex& ix) {
    for (QModelIndex p = ix.parent(); p.isValid(); p = p.parent())
        emit dataChanged(p, p, {Qt::CheckStateRole});
}

void FolderTreeModel::setRecursive(bool on) {
    if (on == m_recursive) return;
    m_recursive = on;
    emitSubtree(index(QDir::rootPath()));
    emit selectionChanged();
}

void FolderTreeModel::setCheckedRoots(const QStringList& roots,
                                      const QStringList& excluded) {
    const auto clean = [](QString p) {
        while (p.size() > 1 && p.endsWith('/')) p.chop(1);
        return p;
    };
    m_checked.clear();
    m_excluded.clear();
    for (const QString& p : roots) {
        const QString c = clean(p);
        if (!c.isEmpty()) check(c);
    }
    for (const QString& p : excluded) {
        const QString c = clean(p);
        if (!c.isEmpty()) uncheck(c);
    }
    emitSubtree(index(QDir::rootPath()));
    emit selectionChanged();
}

void FolderTreeModel::clearChecks() {
    m_checked.clear();
    m_excluded.clear();
    emitSubtree(index(QDir::rootPath()));
    emit selectionChanged();
}

}