aissia/src/services/StorageService.cpp

#include "StorageService.hpp"

#include <spdlog/sinks/stdout_color_sinks.h>
#include <sqlite3.h>
#include <filesystem>
#include <ctime>

namespace fs = std::filesystem;

namespace aissia {

StorageService::StorageService() {
    m_logger = spdlog::get("StorageService");
    if (!m_logger) {
        m_logger = spdlog::stdout_color_mt("StorageService");
    }
}

StorageService::~StorageService() {
    shutdown();
}

bool StorageService::initialize(grove::IIO* io) {
    m_io = io;

    if (m_io) {
        grove::SubscriptionConfig config;
        m_io->subscribe("storage:save_session", config);
        m_io->subscribe("storage:save_app_usage", config);
        m_io->subscribe("storage:save_conversation", config);
        m_io->subscribe("storage:update_metrics", config);
    }

    m_logger->info("StorageService initialized");
    return true;
}

bool StorageService::openDatabase(const std::string& dbPath,
                                   const std::string& journalMode,
                                   int busyTimeoutMs) {
    m_dbPath = dbPath;

    // Ensure directory exists
    fs::path path(dbPath);
    if (path.has_parent_path()) {
        fs::create_directories(path.parent_path());
    }

    int rc = sqlite3_open(dbPath.c_str(), &m_db);
    if (rc != SQLITE_OK) {
        m_logger->error("SQLite open error: {}", sqlite3_errmsg(m_db));
        return false;
    }

    // Set pragmas
    std::string pragmas = "PRAGMA journal_mode=" + journalMode + ";"
                         "PRAGMA busy_timeout=" + std::to_string(busyTimeoutMs) + ";"
                         "PRAGMA foreign_keys=ON;";
    if (!executeSQL(pragmas)) {
        return false;
    }

    if (!initializeSchema()) {
        return false;
    }

    if (!prepareStatements()) {
        return false;
    }

    m_isConnected = true;

    // Publish ready event
    if (m_io) {
        auto event = std::make_unique<grove::JsonDataNode>("ready");
        event->setString("database", dbPath);
        m_io->publish("storage:ready", std::move(event));
    }

    m_logger->info("Database opened: {}", dbPath);
    return true;
}

bool StorageService::initializeSchema() {
    const char* schema = R"SQL(
        CREATE TABLE IF NOT EXISTS work_sessions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            task_name TEXT,
            start_time INTEGER,
            end_time INTEGER,
            duration_minutes INTEGER,
            hyperfocus_detected BOOLEAN DEFAULT 0,
            created_at INTEGER DEFAULT (strftime('%s', 'now'))
        );

        CREATE TABLE IF NOT EXISTS app_usage (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id INTEGER,
            app_name TEXT,
            duration_seconds INTEGER,
            is_productive BOOLEAN,
            created_at INTEGER DEFAULT (strftime('%s', 'now')),
            FOREIGN KEY (session_id) REFERENCES work_sessions(id)
        );

        CREATE TABLE IF NOT EXISTS conversations (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            role TEXT,
            content TEXT,
            provider TEXT,
            model TEXT,
            tokens_used INTEGER,
            created_at INTEGER DEFAULT (strftime('%s', 'now'))
        );

        CREATE TABLE IF NOT EXISTS daily_metrics (
            date TEXT PRIMARY KEY,
            total_focus_minutes INTEGER DEFAULT 0,
            total_breaks INTEGER DEFAULT 0,
            hyperfocus_count INTEGER DEFAULT 0,
            updated_at INTEGER DEFAULT (strftime('%s', 'now'))
        );

        CREATE INDEX IF NOT EXISTS idx_sessions_date ON work_sessions(created_at);
        CREATE INDEX IF NOT EXISTS idx_app_usage_session ON app_usage(session_id);
        CREATE INDEX IF NOT EXISTS idx_conversations_date ON conversations(created_at);
    )SQL";

    return executeSQL(schema);
}

bool StorageService::prepareStatements() {
    int rc;

    // Save session statement
    const char* sqlSession = "INSERT INTO work_sessions "
        "(task_name, start_time, end_time, duration_minutes, hyperfocus_detected) "
        "VALUES (?, ?, ?, ?, ?)";
    rc = sqlite3_prepare_v2(m_db, sqlSession, -1, &m_stmtSaveSession, nullptr);
    if (rc != SQLITE_OK) {
        m_logger->error("Failed to prepare save_session: {}", sqlite3_errmsg(m_db));
        return false;
    }

    // Save app usage statement
    const char* sqlAppUsage = "INSERT INTO app_usage "
        "(session_id, app_name, duration_seconds, is_productive) "
        "VALUES (?, ?, ?, ?)";
    rc = sqlite3_prepare_v2(m_db, sqlAppUsage, -1, &m_stmtSaveAppUsage, nullptr);
    if (rc != SQLITE_OK) {
        m_logger->error("Failed to prepare save_app_usage: {}", sqlite3_errmsg(m_db));
        return false;
    }

    // Save conversation statement
    const char* sqlConv = "INSERT INTO conversations "
        "(role, content, provider, model, tokens_used) "
        "VALUES (?, ?, ?, ?, ?)";
    rc = sqlite3_prepare_v2(m_db, sqlConv, -1, &m_stmtSaveConversation, nullptr);
    if (rc != SQLITE_OK) {
        m_logger->error("Failed to prepare save_conversation: {}", sqlite3_errmsg(m_db));
        return false;
    }

    // Update metrics statement
    const char* sqlMetrics = "INSERT INTO daily_metrics "
        "(date, total_focus_minutes, total_breaks, hyperfocus_count) "
        "VALUES (?, ?, ?, ?) "
        "ON CONFLICT(date) DO UPDATE SET "
        "total_focus_minutes = total_focus_minutes + excluded.total_focus_minutes, "
        "total_breaks = total_breaks + excluded.total_breaks, "
        "hyperfocus_count = hyperfocus_count + excluded.hyperfocus_count, "
        "updated_at = strftime('%s', 'now')";
    rc = sqlite3_prepare_v2(m_db, sqlMetrics, -1, &m_stmtUpdateMetrics, nullptr);
    if (rc != SQLITE_OK) {
        m_logger->error("Failed to prepare update_metrics: {}", sqlite3_errmsg(m_db));
        return false;
    }

    m_logger->debug("Prepared statements created");
    return true;
}

void StorageService::finalizeStatements() {
    if (m_stmtSaveSession) { sqlite3_finalize(m_stmtSaveSession); m_stmtSaveSession = nullptr; }
    if (m_stmtSaveAppUsage) { sqlite3_finalize(m_stmtSaveAppUsage); m_stmtSaveAppUsage = nullptr; }
    if (m_stmtSaveConversation) { sqlite3_finalize(m_stmtSaveConversation); m_stmtSaveConversation = nullptr; }
    if (m_stmtUpdateMetrics) { sqlite3_finalize(m_stmtUpdateMetrics); m_stmtUpdateMetrics = nullptr; }
}

void StorageService::process() {
    processMessages();
}

void StorageService::processMessages() {
    if (!m_io || !m_isConnected) return;

    while (m_io->hasMessages() > 0) {
        auto msg = m_io->pullMessage();

        if (msg.topic == "storage:save_session" && msg.data) {
            handleSaveSession(*msg.data);
        }
        else if (msg.topic == "storage:save_app_usage" && msg.data) {
            handleSaveAppUsage(*msg.data);
        }
        else if (msg.topic == "storage:save_conversation" && msg.data) {
            handleSaveConversation(*msg.data);
        }
        else if (msg.topic == "storage:update_metrics" && msg.data) {
            handleUpdateMetrics(*msg.data);
        }
    }
}

void StorageService::handleSaveSession(const grove::IDataNode& data) {
    std::string taskName = data.getString("taskName", "unknown");
    int durationMinutes = data.getInt("durationMinutes", 0);
    bool hyperfocus = data.getBool("hyperfocus", false);

    std::time_t now = std::time(nullptr);
    std::time_t startTime = now - (durationMinutes * 60);

    sqlite3_reset(m_stmtSaveSession);
    sqlite3_bind_text(m_stmtSaveSession, 1, taskName.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_int64(m_stmtSaveSession, 2, startTime);
    sqlite3_bind_int64(m_stmtSaveSession, 3, now);
    sqlite3_bind_int(m_stmtSaveSession, 4, durationMinutes);
    sqlite3_bind_int(m_stmtSaveSession, 5, hyperfocus ? 1 : 0);

    int rc = sqlite3_step(m_stmtSaveSession);
    if (rc == SQLITE_DONE) {
        m_lastSessionId = static_cast<int>(sqlite3_last_insert_rowid(m_db));
        m_totalQueries++;
        m_logger->debug("Session saved: {} ({}min), id={}", taskName, durationMinutes, m_lastSessionId);

        if (m_io) {
            auto event = std::make_unique<grove::JsonDataNode>("saved");
            event->setInt("sessionId", m_lastSessionId);
            m_io->publish("storage:session_saved", std::move(event));
        }
    } else {
        publishError(sqlite3_errmsg(m_db));
    }
}

void StorageService::handleSaveAppUsage(const grove::IDataNode& data) {
    int sessionId = data.getInt("sessionId", m_lastSessionId);
    std::string appName = data.getString("appName", "");
    int durationSeconds = data.getInt("durationSeconds", 0);
    bool productive = data.getBool("productive", false);

    sqlite3_reset(m_stmtSaveAppUsage);
    sqlite3_bind_int(m_stmtSaveAppUsage, 1, sessionId);
    sqlite3_bind_text(m_stmtSaveAppUsage, 2, appName.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_int(m_stmtSaveAppUsage, 3, durationSeconds);
    sqlite3_bind_int(m_stmtSaveAppUsage, 4, productive ? 1 : 0);

    int rc = sqlite3_step(m_stmtSaveAppUsage);
    if (rc == SQLITE_DONE) {
        m_totalQueries++;
    } else {
        publishError(sqlite3_errmsg(m_db));
    }
}

void StorageService::handleSaveConversation(const grove::IDataNode& data) {
    std::string role = data.getString("role", "");
    std::string content = data.getString("content", "");
    std::string provider = data.getString("provider", "");
    std::string model = data.getString("model", "");
    int tokens = data.getInt("tokens", 0);

    sqlite3_reset(m_stmtSaveConversation);
    sqlite3_bind_text(m_stmtSaveConversation, 1, role.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_text(m_stmtSaveConversation, 2, content.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_text(m_stmtSaveConversation, 3, provider.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_text(m_stmtSaveConversation, 4, model.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_int(m_stmtSaveConversation, 5, tokens);

    int rc = sqlite3_step(m_stmtSaveConversation);
    if (rc == SQLITE_DONE) {
        m_totalQueries++;
    } else {
        publishError(sqlite3_errmsg(m_db));
    }
}

void StorageService::handleUpdateMetrics(const grove::IDataNode& data) {
    int focusMinutes = data.getInt("focusMinutes", 0);
    int breaks = data.getInt("breaks", 0);
    int hyperfocusCount = data.getInt("hyperfocusCount", 0);

    std::time_t now = std::time(nullptr);
    std::tm* tm = std::localtime(&now);
    char dateStr[11];
    std::strftime(dateStr, sizeof(dateStr), "%Y-%m-%d", tm);

    sqlite3_reset(m_stmtUpdateMetrics);
    sqlite3_bind_text(m_stmtUpdateMetrics, 1, dateStr, -1, SQLITE_TRANSIENT);
    sqlite3_bind_int(m_stmtUpdateMetrics, 2, focusMinutes);
    sqlite3_bind_int(m_stmtUpdateMetrics, 3, breaks);
    sqlite3_bind_int(m_stmtUpdateMetrics, 4, hyperfocusCount);

    int rc = sqlite3_step(m_stmtUpdateMetrics);
    if (rc == SQLITE_DONE) {
        m_totalQueries++;
    } else {
        publishError(sqlite3_errmsg(m_db));
    }
}

bool StorageService::executeSQL(const std::string& sql) {
    char* errMsg = nullptr;
    int rc = sqlite3_exec(m_db, sql.c_str(), nullptr, nullptr, &errMsg);

    if (rc != SQLITE_OK) {
        m_logger->error("SQL error: {}", errMsg ? errMsg : "unknown");
        sqlite3_free(errMsg);
        return false;
    }

    m_totalQueries++;
    return true;
}

void StorageService::publishError(const std::string& message) {
    m_logger->error("Storage error: {}", message);
    if (m_io) {
        auto event = std::make_unique<grove::JsonDataNode>("error");
        event->setString("message", message);
        m_io->publish("storage:error", std::move(event));
    }
}

void StorageService::shutdown() {
    finalizeStatements();

    if (m_db) {
        sqlite3_close(m_db);
        m_db = nullptr;
        m_isConnected = false;
    }

    m_logger->info("StorageService shutdown. Total queries: {}", m_totalQueries);
}

} // namespace aissia