#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();
		}

	}