ocr-screenshot-gallery/src/databasemanager.cpp

#include "databasemanager.h"
#include <QSqlQuery>
#include <QSqlError>
#include <QDebug>
#include <QVariant>
#include <QFileInfo>

DatabaseManager::DatabaseManager(QObject *parent)
    : QObject(parent)
    , m_initialized(false)
{
}

DatabaseManager::~DatabaseManager()
{
    if (m_db.isOpen()) {
        m_db.close();
    }
}

bool DatabaseManager::initialize(const QString &dbPath)
{
    // Check if database is already initialized
    if (m_initialized) {
        return true;
    }

    // Check if file exists
    QFileInfo fileInfo(dbPath);
    if (!fileInfo.exists() || !fileInfo.isFile()) {
        qDebug() << "Database file does not exist:" << dbPath;
        return false;
    }

    // Set up database connection
    m_db = QSqlDatabase::addDatabase("QSQLITE");
    m_db.setDatabaseName(dbPath);

    // Open database
    if (!m_db.open()) {
        qDebug() << "Failed to open database:" << m_db.lastError().text();
        return false;
    }

    // Verify required table exists
    QSqlQuery query;
    if (!query.exec("SELECT name FROM sqlite_master WHERE type='table' AND name='ocr_results'")) {
        qDebug() << "Failed to execute query:" << query.lastError().text();
        m_db.close();
        return false;
    }
    
    if (!query.next()) {
        qDebug() << "The required table 'ocr_results' does not exist in the database.";
        m_db.close();
        return false;
    }

    // Verify the table has the required columns
    if (!query.exec("PRAGMA table_info(ocr_results)")) {
        qDebug() << "Failed to get table info:" << query.lastError().text();
        m_db.close();
        return false;
    }
    
    bool hasFullPath = false;
    bool hasOcrText = false;
    
    while (query.next()) {
        QString columnName = query.value(1).toString();
        if (columnName == "full_path") hasFullPath = true;
        if (columnName == "ocr_text") hasOcrText = true;
    }
    
    if (!hasFullPath || !hasOcrText) {
        qDebug() << "Missing required columns in ocr_results table. Need 'full_path' and 'ocr_text'";
        m_db.close();
        return false;
    }

    m_initialized = true;
    qDebug() << "Database initialized successfully.";
    return true;
}

QList<DatabaseManager::ImageItem> DatabaseManager::getAllImages()
{
    QList<ImageItem> images;
    
    if (!m_initialized) {
        qDebug() << "Database not initialized.";
        return images;
    }

    // Verify database is still connected
    if (!m_db.isOpen() && !m_db.open()) {
        qDebug() << "Database connection lost and cannot be reopened:" << m_db.lastError().text();
        m_initialized = false;
        return images;
    }

    QSqlQuery query;
    query.prepare("SELECT full_path, ocr_text FROM ocr_results");
    
    if (!query.exec()) {
        qDebug() << "Failed to fetch images:" << query.lastError().text();
        return images;
    }

    // Check if files exist as we add them
    while (query.next()) {
        ImageItem item;
        item.filePath = query.value(0).toString();
        item.ocrText = query.value(1).toString();
        
        // Only add images that have a non-empty path
        if (!item.filePath.isEmpty()) {
            images.append(item);
        }
    }

    return images;
}

QList<DatabaseManager::ImageItem> DatabaseManager::searchImages(const QString &searchText)
{
    QList<ImageItem> images;
    
    if (!m_initialized) {
        qDebug() << "Database not initialized.";
        return images;
    }

    // Verify database is still connected
    if (!m_db.isOpen() && !m_db.open()) {
        qDebug() << "Database connection lost and cannot be reopened:" << m_db.lastError().text();
        m_initialized = false;
        return images;
    }

    // If search text is empty, return all images
    if (searchText.isEmpty()) {
        return getAllImages();
    }

    QSqlQuery query;
    // Use LIKE query with wildcards for flexible searching
    query.prepare("SELECT full_path, ocr_text FROM ocr_results WHERE ocr_text LIKE :search");
    
    // Ensure search text is properly sanitized
    QString sanitizedSearch = searchText;
    sanitizedSearch.replace('\'', "''"); // Escape single quotes
    
    query.bindValue(":search", "%" + sanitizedSearch + "%");
    
    if (!query.exec()) {
        qDebug() << "Failed to search images:" << query.lastError().text();
        return images;
    }

    while (query.next()) {
        ImageItem item;
        item.filePath = query.value(0).toString();
        item.ocrText = query.value(1).toString();
        
        // Only add images that have a non-empty path
        if (!item.filePath.isEmpty()) {
            images.append(item);
        }
    }

    return images;
}