z-bar-qt/Plugins/ZShell/Models/wallpaperdb.cpp

#include "wallpaperdb.hpp"
#include "wallpaperutil.hpp"

#include <QtQuick/qquickimageprovider.h>
#include <qbytearray.h>
#include <qdiriterator.h>
#include <qelapsedtimer.h>
#include <qfuturewatcher.h>
#include <qimage.h>
#include <qimagereader.h>
#include <qjsondocument.h>
#include <qjsonobject.h>
#include <qlogging.h>
#include <qqmlengine.h>
#include <qsqldatabase.h>
#include <qsqlerror.h>
#include <qsqlquery.h>
#include <qtconcurrentrun.h>
#include <quuid.h>

#include <algorithm>
#include <atomic>
#include <mutex>

Q_LOGGING_CATEGORY(dbLog, "ZShell.Internal.wallpaper.db", QtWarningMsg)

namespace ZShell::models {

class WallpaperImageProvider : public QQuickImageProvider {
	public:
	explicit WallpaperImageProvider(const QString& dbPath)
		: QQuickImageProvider(QQuickImageProvider::Image), m_dbPath(dbPath) {}

	void setDbPath(const QString& dbPath) {
		std::lock_guard<std::mutex> lock(m_mutex);
		if (m_dbPath == dbPath) {
			return;
		}

		m_dbPath = dbPath;
		m_generation.fetch_add(1, std::memory_order_relaxed);
	}

	QImage requestImage(
		const QString& id, QSize* size, const QSize& requestedSize) override {
		const QString path =
			QString::fromUtf8(QByteArray::fromHex(id.toUtf8()));
		if (path.isEmpty() || !openThreadDb(currentDbPath())) {
			return {};
		}

		QSqlQuery query(QSqlDatabase::database(m_threadUuid));
		query.prepare("SELECT cache FROM wallpapers WHERE path = ?");
		query.addBindValue(path);

		if (!query.exec() || !query.next()) {
			return {};
		}

		const QByteArray bytes = query.value(0).toByteArray();
		if (bytes.isEmpty()) {
			return {};
		}

		QImage image = QImage::fromData(bytes, "JPEG");
		if (image.isNull()) {
			return {};
		}

		if (size) {
			*size = image.size();
		}

		if (requestedSize.isValid() && !requestedSize.isEmpty() &&
			(image.width() > requestedSize.width() ||
			 image.height() > requestedSize.height())) {
			image = image.scaled(
				requestedSize, Qt::KeepAspectRatio, Qt::SmoothTransformation);
		}

		return image;
	}

	private:
	QString currentDbPath() const {
		std::lock_guard<std::mutex> lock(m_mutex);
		return m_dbPath;
	}

	bool openThreadDb(const QString& dbPath) {
		const quint64 generation = m_generation.load(std::memory_order_relaxed);

		if (!m_threadUuid.isEmpty() && m_threadGeneration == generation) {
			return true;
		}

		if (!m_threadUuid.isEmpty()) {
			{
				auto db = QSqlDatabase::database(m_threadUuid, false);
				if (db.isValid()) {
					db.close();
				}
			}

			QSqlDatabase::removeDatabase(m_threadUuid);
			m_threadUuid.clear();
		}

		if (dbPath.isEmpty()) {
			return false;
		}

		const QString uuid = QUuid::createUuid().toString(QUuid::WithoutBraces);

		auto db = QSqlDatabase::addDatabase("QSQLITE", uuid);
		db.setConnectOptions("QSQLITE_OPEN_READONLY");
		db.setDatabaseName(dbPath);

		if (!db.open()) {
			return false;
		}

		QSqlQuery query(db);
		query.exec("PRAGMA journal_mode=WAL");

		m_threadUuid = uuid;
		m_threadGeneration = generation;
		return true;
	}

	mutable std::mutex m_mutex;
	QString m_dbPath;
	std::atomic<quint64> m_generation = 0;

	thread_local static QString m_threadUuid;
	thread_local static quint64 m_threadGeneration;
};

thread_local QString WallpaperImageProvider::m_threadUuid;
thread_local quint64 WallpaperImageProvider::m_threadGeneration;

WallpaperEntry::WallpaperEntry(
	const QString& path, const QString& relativePath, QObject* parent)
	: QObject(parent)
	, m_fileInfo(path)
	, m_path(path)
	, m_relativePath(relativePath) {}

QString WallpaperEntry::path() const {
	return m_path;
}

QString WallpaperEntry::relativePath() const {
	return m_relativePath;
}

QString WallpaperEntry::name() const {
	return m_fileInfo.fileName();
}

QString WallpaperEntry::baseName() const {
	return m_fileInfo.baseName();
}

QString WallpaperEntry::parentDir() const {
	return m_fileInfo.absolutePath();
}

QString WallpaperEntry::suffix() const {
	return m_fileInfo.completeSuffix();
}

qint64 WallpaperEntry::size() const {
	return m_fileInfo.size();
}

bool WallpaperEntry::isImage() const {
	return !hasVideoExtension(m_path);
}

bool WallpaperEntry::isVideo() const {
	return hasVideoExtension(m_path);
}

bool WallpaperEntry::hardware() const {
	return m_hardware;
}

bool WallpaperEntry::cached() const {
	return !m_thumbnailPath.isEmpty();
}

QString WallpaperEntry::thumbnailPath() const {
	if (!m_thumbnailPath.isEmpty()) {
		return m_thumbnailPath;
	}

	return isImage() ? m_path : QString();
}

void WallpaperEntry::setRelativePath(const QString& path) {
	if (m_relativePath == path) {
		return;
	}

	m_relativePath = path;
	emit relativePathChanged();
}

void WallpaperEntry::setHardware(bool hardware) {
	if (m_hardware == hardware) {
		return;
	}

	m_hardware = hardware;
	emit hardwareChanged();
}

void WallpaperEntry::setThumbnailPath(const QString& path) {
	if (m_thumbnailPath == path) {
		return;
	}

	m_thumbnailPath = path;
	emit thumbnailPathChanged();
}

WallpaperDb::WallpaperDb(QObject* parent) : QObject(parent) {
	m_uuid = QUuid::createUuid().toString(QUuid::WithoutBraces);
	m_dbPath = QDir::homePath() + "/.local/state/zshell/wallpapers.sqlite";

	m_queueTimer.setInterval(300);
	m_queueTimer.setTimerType(Qt::CoarseTimer);
	QObject::connect(
		&m_queueTimer, &QTimer::timeout, this, &WallpaperDb::processQueue);

	m_watchTimer.setInterval(400);
	m_watchTimer.setSingleShot(true);
	QObject::connect(
		&m_watcher, &QFileSystemWatcher::directoryChanged, this, [this]() {
			m_watchTimer.start();
		});
	QObject::connect(
		&m_watchTimer, &QTimer::timeout, this, &WallpaperDb::scanDir);
}

WallpaperDb::~WallpaperDb() {
	if (QSqlDatabase::contains(m_uuid)) {
		{
			auto db = QSqlDatabase::database(m_uuid, false);
			if (db.isValid()) {
				db.close();
			}
		}

		QSqlDatabase::removeDatabase(m_uuid);
	}
}

QString WallpaperDb::path() const {
	return m_path;
}

void WallpaperDb::setPath(const QString& path) {
	if (m_path == path) {
		return;
	}

	m_path = path;
	m_dir = QDir(path);
	emit pathChanged();

	for (const auto& dir : m_watcher.directories()) {
		m_watcher.removePath(dir);
	}

	scanDir();
}

QString WallpaperDb::dbPath() const {
	return m_dbPath;
}

void WallpaperDb::setDbPath(const QString& path) {
	if (m_dbPath == path) {
		return;
	}

	m_dbPath = path;
	emit dbPathChanged();

	if (!m_path.isEmpty()) {
		scanDir();
	}
}

bool WallpaperDb::requireHardwareCodec() const {
	return m_requireHardwareCodec;
}

void WallpaperDb::setRequireHardwareCodec(bool requireHardwareCodec) {
	if (m_requireHardwareCodec == requireHardwareCodec) {
		return;
	}

	m_requireHardwareCodec = requireHardwareCodec;
	emit requireHardwareCodecChanged();

	rebuildEntries();
}

QQmlListProperty<WallpaperEntry> WallpaperDb::entries() {
	return QQmlListProperty<WallpaperEntry>(this, &m_entries);
}

QString WallpaperDb::thumbnailUrl(const QString& path) {
	return QStringLiteral("image://wallpapers/") +
		   QString::fromUtf8(path.toUtf8().toHex());
}

QSize WallpaperDb::cacheSize() const {
	return m_cacheSize;
}

void WallpaperDb::setCacheSize(const QSize& cacheSize) {
	if (m_cacheSize == cacheSize) {
		return;
	}

	const QSize prev = m_cacheSize;
	const bool hadPrevious = !prev.isEmpty();

	m_cacheSize = cacheSize;
	emit cacheSizeChanged();

	if (hadPrevious && qint64(cacheSize.width()) * cacheSize.height() >
						   qint64(prev.width()) * prev.height()) {
		const auto db = QSqlDatabase::database(m_uuid, false);

		for (auto* entry : std::as_const(m_byPath)) {
			if (!entry->cached()) {
				continue;
			}

			if (db.isOpen()) {
				QSqlQuery query(db);
				query.prepare(
					"UPDATE wallpapers SET cache = NULL, scheme_thumb = NULL,"
					" poster = NULL WHERE path = ?");
				query.addBindValue(entry->path());
				query.exec();
			}

			entry->setThumbnailPath(QString());
			queueMetadata(entry, entry->isVideo());
		}

		rebuildEntries();
	}

	if (!m_queue.isEmpty() && !m_queueTimer.isActive()) {
		m_queueTimer.start();
	}
}

void WallpaperDb::registerProvider() {
	if (m_provider != nullptr) {
		m_provider->setDbPath(m_dbPath);
		return;
	}

	auto* engine = qmlEngine(this);
	if (engine == nullptr) {
		return;
	}

	m_provider = new WallpaperImageProvider(m_dbPath);
	engine->addImageProvider(QStringLiteral("wallpapers"), m_provider);
}

void WallpaperDb::openDb() {
	if (m_dbPath.isEmpty()) {
		return;
	}

	if (QSqlDatabase::contains(m_uuid)) {
		{
			auto db = QSqlDatabase::database(m_uuid, false);
			if (db.isOpen() && db.databaseName() == m_dbPath) {
				registerProvider();
				return;
			}

			if (db.isValid()) {
				db.close();
			}
		}

		QSqlDatabase::removeDatabase(m_uuid);
	}

	QDir().mkpath(QFileInfo(m_dbPath).absolutePath());

	auto db = QSqlDatabase::addDatabase("QSQLITE", m_uuid);
	db.setDatabaseName(m_dbPath);

	if (!db.open()) {
		qWarning(dbLog) << "WallpaperDb: could not open database" << m_dbPath
						<< db.lastError();
		return;
	}

	QSqlQuery query(db);
	query.exec("PRAGMA journal_mode=WAL");
	query.exec(
		"CREATE TABLE IF NOT EXISTS wallpapers("
		"path TEXT PRIMARY KEY, kind TEXT, codec TEXT, hardware INTEGER,"
		" size INTEGER, mtime INTEGER, cache BLOB, scheme_thumb BLOB,"
		" poster BLOB, crop TEXT)");

	query.exec("SELECT cache FROM wallpapers LIMIT 1");
	if (query.lastError().isValid()) {
		query.exec("DROP TABLE IF EXISTS wallpapers");
		query.exec(
			"CREATE TABLE wallpapers("
			"path TEXT PRIMARY KEY, kind TEXT, codec TEXT, hardware INTEGER,"
			" size INTEGER, mtime INTEGER, cache BLOB, scheme_thumb BLOB,"
			" poster BLOB, crop TEXT)");
	} else {
		query.exec("ALTER TABLE wallpapers ADD COLUMN scheme_thumb BLOB");
		query.exec("ALTER TABLE wallpapers ADD COLUMN poster BLOB");
		query.exec("ALTER TABLE wallpapers ADD COLUMN crop TEXT");
	}

	registerProvider();
}

void WallpaperDb::scanDir() {
	openDb();

	watchDir(m_dir.absolutePath());

	if (!m_dir.exists()) {
		if (!m_entries.isEmpty()) {
			m_entries.clear();
			emit entriesChanged();
		}

		return;
	}

	auto* watcher = new QFutureWatcher<ScanResult>(this);
	QObject::connect(
		watcher, &QFutureWatcher<ScanResult>::finished, this, [this, watcher]() {
			applyScan(watcher->result());
			watcher->deleteLater();
		});

	const QDir dir = m_dir;

	QHash<QString, QPair<qint64, bool>> knownCodecs;
	for (const auto* entry : std::as_const(m_byPath)) {
		if (entry->isVideo()) {
			knownCodecs.insert(
				entry->path(), qMakePair(entry->size(), entry->hardware()));
		}
	}

	watcher->setFuture(QtConcurrent::run([dir, knownCodecs]() {
		ScanResult result;
		result.dir = dir;

		QDirIterator it(dir.absolutePath(), QDirIterator::Subdirectories);
		while (it.hasNext()) {
			const QString path = it.next();
			const QFileInfo info(path);

			if (info.isDir()) {
				continue;
			}

			if (hasVideoExtension(path)) {
				FileEntry file;
				file.info = info;

				const auto known = knownCodecs.constFind(path);
				if (known != knownCodecs.constEnd() &&
					known.value().first == info.size()) {
					file.hardware = known.value().second;
				} else {
					file.hardware = probeHardwareCodec(path);
				}

				result.files.append(file);
			} else {
				QImageReader reader(path);
				if (reader.canRead()) {
					FileEntry file;
					file.info = info;
					result.files.append(file);
				}
			}
		}

		return result;
	}));
}

void WallpaperDb::watchDir(const QString& dir) {
	if (dir.isEmpty() || m_watcher.directories().contains(dir)) {
		return;
	}

	if (!QDir(dir).exists()) {
		return;
	}

	m_watcher.addPath(dir);
}

void WallpaperDb::applyScan(const ScanResult& result) {
	qCInfo(
		dbLog,
		"WallpaperDb: scan of %s found %d files (cache target %dx%d)",
		m_dir.absolutePath().toUtf8().constData(),
		int(result.files.size()),
		m_cacheSize.width(),
		m_cacheSize.height());

	const auto db = QSqlDatabase::database(m_uuid, false);

	QSet<QString> seen;
	for (const auto& file : result.files) {
		seen.insert(file.info.absoluteFilePath());
	}

	for (auto it = m_byPath.begin(); it != m_byPath.end();) {
		if (!seen.contains(it.key())) {
			if (db.isOpen()) {
				QSqlQuery query(db);
				query.prepare("DELETE FROM wallpapers WHERE path = ?");
				query.addBindValue(it.key());
				query.exec();
			}

			delete it.value();
			m_crops.remove(it.key());
			it = m_byPath.erase(it);
		} else {
			++it;
		}
	}

	for (const auto& file : result.files) {
		const QFileInfo& info = file.info;
		const QString path = info.absoluteFilePath();

		const qint64 size = info.size();
		const qint64 mtime = info.lastModified().toMSecsSinceEpoch() / 1000;
		const bool isVideo = hasVideoExtension(path);

		watchDir(info.absolutePath());

		const auto existing = m_byPath.constFind(path);
		if (existing != m_byPath.constEnd()) {
			if (existing.value()->size() == size) {
				existing.value()->setHardware(file.hardware);
				continue;
			}
		}

		WallpaperEntry* entry = m_byPath.value(path);
		if (!entry) {
			entry =
				new WallpaperEntry(path, m_dir.relativeFilePath(path), this);
			m_byPath.insert(path, entry);
		}

		entry->setHardware(file.hardware);

		bool reused = false;
		if (db.isOpen()) {
			QSqlQuery query(db);
			query.prepare(
				"SELECT cache, scheme_thumb, poster, size, mtime, crop"
				" FROM wallpapers WHERE path = ?");
			query.addBindValue(path);
			query.exec();

			if (query.next()) {
				const auto cache = query.value(0).toByteArray();
				const auto scheme = query.value(1).toByteArray();
				const auto poster = query.value(2).toByteArray();
				const auto storedSize = query.value(3).toLongLong();
				const auto storedMtime = query.value(4).toLongLong();
				loadCrop(path, query.value(5).toString());

				const bool posterOk = !isVideo || !poster.isEmpty();

				if (storedSize == size && storedMtime == mtime &&
					!cache.isEmpty()) {
					reused = true;
					entry->setThumbnailPath(thumbnailUrl(path));

					if (scheme.isEmpty() || !posterOk) {
						queueMetadata(entry, isVideo, cache, poster);
					}
				}
			}
		}

		if (!reused) {
			if (db.isOpen()) {
				QSqlQuery query(db);
				query.prepare(
					"INSERT INTO wallpapers"
					" (path, kind, hardware, size, mtime, cache)"
					" VALUES (?, ?, ?, ?, ?, NULL)"
					" ON CONFLICT(path) DO UPDATE SET kind = excluded.kind,"
					" hardware = excluded.hardware, size = excluded.size,"
					" mtime = excluded.mtime, cache = NULL,"
					" scheme_thumb = NULL, poster = NULL");
				query.addBindValue(path);
				query.addBindValue(isVideo ? "video" : "image");
				query.addBindValue(
					isVideo ? (file.hardware ? 1 : 0) : QVariant());
				query.addBindValue(size);
				query.addBindValue(mtime);
				if (!query.exec()) {
					qWarning(dbLog)
						<< "WallpaperDb: could not store the row for" << path
						<< query.lastError();
				}
			}

			queueMetadata(entry, isVideo);
		}
	}

	rebuildEntries();
}

void WallpaperDb::queueMetadata(
	WallpaperEntry* entry,
	bool isVideo,
	const QByteArray& existingCache,
	const QByteArray& existingPoster) {
	qCInfo(
		dbLog,
		"WallpaperDb: queued cache job for %s (%s%s)",
		entry->path().toUtf8().constData(),
		isVideo ? "video" : "image",
		existingCache.isEmpty() ? "" : ", existing blob");

	m_queue.append(
		MetadataJob{entry, entry->path(), isVideo, existingCache, existingPoster});

	if (!m_queueTimer.isActive()) {
		m_queueTimer.start();
	}
}

void WallpaperDb::processQueue() {
	if (m_jobRunning) {
		return;
	}

	if (m_queue.isEmpty()) {
		m_queueTimer.stop();
		return;
	}

	if (m_cacheSize.isEmpty()) {
		return;
	}

	const MetadataJob job = m_queue.takeFirst();
	m_jobRunning = true;

	qCInfo(
		dbLog,
		"WallpaperDb: caching %s (%s, %d jobs left in queue)",
		job.path.toUtf8().constData(),
		job.isVideo ? "video" : "image",
		int(m_queue.size()));

	runMetadataJob(job);
}

void WallpaperDb::runMetadataJob(const MetadataJob& job) {
	struct CacheResult {
		QString path;
		QByteArray cache;
		QByteArray poster;
		QByteArray scheme;
	};

	const QSize cacheSize = m_cacheSize;
	const bool isVideoJob = job.isVideo;

	QElapsedTimer clock;
	clock.start();

	auto* watcher = new QFutureWatcher<CacheResult>(this);
	QObject::connect(
		watcher,
		&QFutureWatcher<CacheResult>::finished,
		this,
		[this, watcher, clock, isVideoJob]() {
			const auto result = watcher->result();

			qCInfo(
				dbLog,
				"WallpaperDb: %s produced %d bytes of cache in %dms",
				result.path.toUtf8().constData(),
				int(result.cache.size()),
				int(clock.elapsed()));

			qCInfo(
				dbLog,
				"WallpaperDb: %s poster %d bytes",
				result.path.toUtf8().constData(),
				int(result.poster.size()));

			if (isVideoJob && result.poster.isEmpty()) {
				qWarning(dbLog)
					<< "WallpaperDb: no poster produced for" << result.path;
			}

			if (QSqlDatabase::contains(m_uuid)) {
				const auto db = QSqlDatabase::database(m_uuid, false);
				if (db.isOpen()) {
					QSqlQuery query(db);
					query.prepare(
						"UPDATE wallpapers SET cache = ?, scheme_thumb = ?,"
						" poster = ? WHERE path = ?");
					query.addBindValue(result.cache);
					query.addBindValue(result.scheme);
					query.addBindValue(result.poster);
					query.addBindValue(result.path);
					if (!query.exec()) {
						qWarning(dbLog)
							<< "WallpaperDb: could not store the cache for"
							<< result.path << query.lastError();
					}
				}
			}

			if (!result.cache.isEmpty()) {
				if (auto* entry = m_byPath.value(result.path)) {
					entry->setThumbnailPath(thumbnailUrl(result.path));
				}

				emit cacheStored(result.path);
			}

			m_jobRunning = false;
			rebuildEntries();
			watcher->deleteLater();
		});

	watcher->setFuture(QtConcurrent::run([job, cacheSize]() {
		CacheResult result;
		result.path = job.path;

		QByteArray cache = job.existingCache;
		if (cache.isEmpty()) {
			cache = job.isVideo ? videoCacheBytes(job.path, cacheSize)
								: imageCacheBytes(job.path, cacheSize);
		}

		QByteArray poster = job.existingPoster;
		if (poster.isEmpty()) {
			poster =
				job.isVideo ? videoPosterBytes(job.path, cacheSize) : cache;
		}

		result.cache = cache;
		result.poster = poster;
		result.scheme = cache.isEmpty() ? QByteArray()
										: schemeThumbBytes(cache, 128);

		return result;
	}));
}

void WallpaperDb::rebuildEntries() {
	QList<WallpaperEntry*> next;
	next.reserve(m_byPath.size());

	for (auto* entry : std::as_const(m_byPath)) {
		if (entryAllowed(entry)) {
			next.append(entry);
		}
	}

	std::sort(
		next.begin(),
		next.end(),
		[](const WallpaperEntry* a, const WallpaperEntry* b) {
			return a->relativePath() < b->relativePath();
		});

	if (next == m_entries) {
		return;
	}

	m_entries = next;
	emit entriesChanged();
}

bool WallpaperDb::entryAllowed(const WallpaperEntry* entry) const {
	if (!entry->isVideo()) {
		return true;
	}

	if (!m_requireHardwareCodec) {
		return true;
	}

	return entry->hardware();
}

QVariantMap WallpaperDb::crop(const QString& path, const QString& screen) const {
	const auto pathCrops = m_crops.constFind(path);
	if (pathCrops == m_crops.constEnd()) return {};
	return pathCrops.value().value(screen).toMap();
}

void WallpaperDb::setCrop(
	const QString& path, const QString& screen, const QVariantMap& crop) {
	if (path.isEmpty() || screen.isEmpty()) return;

	QVariantMap& pathCrops = m_crops[path];
	pathCrops[screen] = crop;
	storeCropJson(path);
	emit cropsChanged(path);
}

void WallpaperDb::loadCrop(const QString& path, const QString& json) {
	if (json.isEmpty()) {
		m_crops.remove(path);
		return;
	}

	const auto doc = QJsonDocument::fromJson(json.toUtf8());
	if (!doc.isObject()) {
		m_crops.remove(path);
		return;
	}

	m_crops[path] = doc.object().toVariantMap();
}

void WallpaperDb::storeCropJson(const QString& path) {
	const auto db = QSqlDatabase::database(m_uuid, false);
	if (!db.isOpen()) return;

	const auto pathCrops = m_crops.value(path);
	const QByteArray json =
		pathCrops.isEmpty()
			? QByteArray()
			: QJsonDocument(QJsonObject::fromVariantMap(pathCrops))
				  .toJson(QJsonDocument::Compact);

	QSqlQuery query(db);
	query.prepare("UPDATE wallpapers SET crop = ? WHERE path = ?");
	query.addBindValue(json.isEmpty() ? QVariant() : QVariant(json));
	query.addBindValue(path);
	if (!query.exec()) {
		qWarning(dbLog) << "WallpaperDb: could not store the crop for" << path
						<< query.lastError();
	}
}

} // namespace ZShell::models