smartbotic-database/service/src/database_service.cpp

#include "database_service.hpp"

#include <grpcpp/grpcpp.h>
#include <spdlog/spdlog.h>

#include <fstream>

namespace smartbotic::database {

DatabaseService::DatabaseService(Config config)
    : config_(std::move(config))
{
}

DatabaseService::~DatabaseService() {
    stop();
}

bool DatabaseService::initialize() {
    spdlog::info("Initializing database service (node: {})", config_.nodeId);

    try {
        // Create data directory if needed
        std::error_code ec;
        std::filesystem::create_directories(config_.dataDirectory, ec);
        if (ec) {
            spdlog::error("Failed to create data directory: {}", ec.message());
            return false;
        }

        setupComponents();

        // Initialize encryption
        if (config_.encryptionEnabled) {
            if (!encryption_->initialize()) {
                spdlog::error("Failed to initialize encryption");
                return false;
            }
        }

        // Recover from persistence
        persistence_->recover(*store_);

        // Set replication sequence after recovery (WAL sequence is now known)
        replication_->setSequence(persistence_->currentWalSequence());

        // Run migrations if enabled
        if (config_.migrations.enabled && !config_.migrations.directory.empty()) {
            if (!runMigrations()) {
                if (config_.migrations.failOnError) {
                    spdlog::error("Failed to run migrations");
                    return false;
                }
                spdlog::warn("Some migrations failed, continuing anyway");
            }
        }

        spdlog::info("Database service initialized successfully");
        return true;

    } catch (const std::exception& e) {
        spdlog::error("Failed to initialize database service: {}", e.what());
        return false;
    }
}

bool DatabaseService::runMigrations() {
    if (!config_.migrations.enabled) {
        return true;
    }

    MigrationRunner::Config migrationConfig;
    migrationConfig.directory = config_.migrations.directory;
    migrationConfig.autoApply = config_.migrations.autoApply;
    migrationConfig.failOnError = config_.migrations.failOnError;

    migrationRunner_ = std::make_unique<MigrationRunner>(*store_, migrationConfig);
    return migrationRunner_->runMigrations();
}

void DatabaseService::start() {
    if (running_.exchange(true)) {
        return;
    }

    spdlog::info("Starting database service on {}:{}", config_.bindAddress, config_.rpcPort);

    // Start components
    store_->start();
    persistence_->start();
    events_->start();
    files_->start();

    if (config_.replicationEnabled) {
        replication_->start();
        // Set initial local collections for discovery
        replication_->setLocalCollections(store_->listCollections());
    }

    // Start gRPC server
    startGrpcServer();

    spdlog::info("Database service started");
}

void DatabaseService::stop() {
    if (!running_.exchange(false)) {
        return;
    }

    spdlog::info("Stopping database service...");

    // Stop gRPC server first
    stopGrpcServer();

    // Stop components in reverse order
    if (replication_) {
        replication_->stop();
    }
    if (files_) {
        files_->stop();
    }
    if (events_) {
        events_->stop();
    }
    if (persistence_) {
        persistence_->stop();
    }
    if (store_) {
        store_->stop();
    }

    spdlog::info("Database service stopped");
}

void DatabaseService::wait() {
    if (serverThread_.joinable()) {
        serverThread_.join();
    }
}

void DatabaseService::signalStop() {
    stopRequested_ = true;
    stop();
}

nlohmann::json DatabaseService::getStats() const {
    nlohmann::json stats;

    if (store_) {
        auto storeStats = store_->getStats();
        stats["documents"] = storeStats.totalDocuments;
        stats["collections"] = storeStats.totalCollections;
        stats["memory_bytes"] = storeStats.estimatedMemoryBytes;
        stats["inserts"] = storeStats.insertCount;
        stats["updates"] = storeStats.updateCount;
        stats["deletes"] = storeStats.deleteCount;
        stats["queries"] = storeStats.queryCount;
    }

    if (persistence_) {
        auto persistStats = persistence_->getStats();
        stats["wal_sequence"] = persistStats.walSequence;
        stats["wal_size_bytes"] = persistStats.walSizeBytes;
        stats["snapshot_count"] = persistStats.snapshotCount;
        stats["last_snapshot_sequence"] = persistStats.lastSnapshotSequence;
    }

    if (events_) {
        stats["subscriptions"] = events_->subscriptionCount();
    }

    return stats;
}

DatabaseService::Config DatabaseService::loadConfig(const std::filesystem::path& configPath) {
    std::ifstream file(configPath);
    if (!file) {
        throw std::runtime_error("Failed to open config file: " + configPath.string());
    }

    nlohmann::json json;
    file >> json;

    Config config;

    // Check for both "storage" and "database" keys for backward compatibility
    nlohmann::json* dbConfig = nullptr;
    if (json.contains("database")) {
        dbConfig = &json["database"];
    } else if (json.contains("storage")) {
        dbConfig = &json["storage"];
    }

    if (dbConfig) {
        auto& db = *dbConfig;

        config.nodeId = db.value("node_id", config.nodeId);
        config.bindAddress = db.value("bind_address", config.bindAddress);
        config.rpcPort = db.value("rpc_port", config.rpcPort);

        // Expand environment variables in data directory
        std::string dataDir = db.value("data_directory", "");
        if (dataDir.find("${HOME}") != std::string::npos) {
            const char* home = std::getenv("HOME");
            if (home) {
                size_t pos = dataDir.find("${HOME}");
                dataDir.replace(pos, 7, home);
            }
        }
        config.dataDirectory = dataDir;

        // Migrations settings
        if (db.contains("migrations")) {
            auto& migrations = db["migrations"];
            config.migrations.enabled = migrations.value("enabled", config.migrations.enabled);
            config.migrations.autoApply = migrations.value("auto_apply", config.migrations.autoApply);
            config.migrations.failOnError = migrations.value("fail_on_error", config.migrations.failOnError);

            std::string migDir = migrations.value("directory", "");
            if (migDir.find("${HOME}") != std::string::npos) {
                const char* home = std::getenv("HOME");
                if (home) {
                    size_t pos = migDir.find("${HOME}");
                    migDir.replace(pos, 7, home);
                }
            }
            config.migrations.directory = migDir;
        }

        // Memory eviction settings
        if (db.contains("memory")) {
            auto& memory = db["memory"];
            config.maxMemoryMb = memory.value("max_memory_mb", config.maxMemoryMb);
            config.evictionThresholdPercent = memory.value("eviction_threshold_percent", config.evictionThresholdPercent);
            config.evictionTargetPercent = memory.value("eviction_target_percent", config.evictionTargetPercent);
            config.evictionCheckIntervalMs = memory.value("eviction_check_interval_ms", config.evictionCheckIntervalMs);
        }

        // Persistence settings
        if (db.contains("persistence")) {
            auto& persistence = db["persistence"];
            config.walSyncIntervalMs = persistence.value("wal_sync_interval_ms", config.walSyncIntervalMs);
            config.snapshotIntervalSec = persistence.value("snapshot_interval_sec", config.snapshotIntervalSec);
            config.compressionEnabled = persistence.value("compression", "lz4") != "none";
        }

        // File settings
        if (db.contains("files")) {
            auto& files = db["files"];
            config.maxFileSizeMb = files.value("max_file_size_mb", config.maxFileSizeMb);
            config.fileCleanupIntervalSec = files.value("cleanup_orphans_interval_sec", config.fileCleanupIntervalSec);
            if (files.contains("allowed_types")) {
                for (const auto& type : files["allowed_types"]) {
                    config.allowedFileTypes.push_back(type.get<std::string>());
                }
            }
        }

        // Encryption settings
        if (db.contains("encryption")) {
            auto& encryption = db["encryption"];
            config.encryptionEnabled = encryption.value("enabled", config.encryptionEnabled);
            config.autoGenerateKey = encryption.value("auto_generate_key", config.autoGenerateKey);

            std::string keyFile = encryption.value("key_file", "");
            if (keyFile.find("${HOME}") != std::string::npos) {
                const char* home = std::getenv("HOME");
                if (home) {
                    size_t pos = keyFile.find("${HOME}");
                    keyFile.replace(pos, 7, home);
                }
            }
            config.keyFilePath = keyFile;
        }

        // Replication settings
        if (db.contains("replication")) {
            auto& replication = db["replication"];
            config.replicationEnabled = replication.value("enabled", config.replicationEnabled);
            config.conflictResolution = replication.value("conflict_resolution", config.conflictResolution);
            if (replication.contains("peers")) {
                for (const auto& peer : replication["peers"]) {
                    config.peerAddresses.push_back(peer.get<std::string>());
                }
            }
        }
    }

    return config;
}

void DatabaseService::setupComponents() {
    // Create memory store with eviction config
    MemoryStore::Config storeConfig;
    storeConfig.nodeId = config_.nodeId;
    storeConfig.maxMemoryBytes = config_.maxMemoryMb * 1024ULL * 1024ULL;
    storeConfig.evictionThresholdPercent = config_.evictionThresholdPercent;
    storeConfig.evictionTargetPercent = config_.evictionTargetPercent;
    storeConfig.evictionCheckIntervalMs = config_.evictionCheckIntervalMs;
    store_ = std::make_unique<MemoryStore>(storeConfig);

    // Create persistence manager
    PersistenceManager::Config persistConfig;
    persistConfig.dataDir = config_.dataDirectory;
    persistConfig.walSyncIntervalMs = config_.walSyncIntervalMs;
    persistConfig.snapshotIntervalSec = config_.snapshotIntervalSec;
    persistConfig.compressionEnabled = config_.compressionEnabled;
    persistence_ = std::make_unique<PersistenceManager>(persistConfig);

    // Connect store callbacks to persistence - uses setPersistCallback for WAL logging
    store_->setPersistCallback([this](const std::string& collection, const std::string& id,
                                       const std::optional<Document>& doc, EventType eventType) {
        // Log to WAL based on event type
        switch (eventType) {
            case EventType::INSERT:
                if (doc) persistence_->logInsert(collection, *doc);
                break;
            case EventType::UPDATE:
                if (doc) persistence_->logUpdate(collection, *doc);
                break;
            case EventType::DELETE:
                persistence_->logDelete(collection, id);
                break;
            default:
                break;
        }

        // Queue for replication broadcast
        if (replication_ && config_.replicationEnabled) {
            databasepb::ReplicationEntry entry;
            entry.set_collection(collection);
            entry.set_document_id(id);
            entry.set_global_timestamp(
                std::chrono::duration_cast<std::chrono::milliseconds>(
                    std::chrono::system_clock::now().time_since_epoch()
                ).count());

            switch (eventType) {
                case EventType::INSERT:
                    entry.set_op(databasepb::OP_INSERT);
                    if (doc) entry.set_data(doc->data.dump());
                    break;
                case EventType::UPDATE:
                    entry.set_op(databasepb::OP_UPDATE);
                    if (doc) entry.set_data(doc->data.dump());
                    break;
                case EventType::DELETE:
                    entry.set_op(databasepb::OP_DELETE);
                    break;
                default:
                    break;
            }

            replication_->queueForReplication(entry);
        }

        // Publish event
        if (events_) {
            DatabaseEvent event;
            event.type = eventType;
            event.collection = collection;
            event.documentId = id;
            event.timestamp = std::chrono::duration_cast<std::chrono::milliseconds>(
                std::chrono::system_clock::now().time_since_epoch()
            ).count();
            event.nodeId = config_.nodeId;
            if (doc) {
                event.data = doc->data;
            }
            events_->publish(event);
        }
    });

    // Set up document load callback for LRU eviction recovery
    store_->setDocumentLoadCallback([this](const std::string& collection,
                                           const std::string& id,
                                           uint64_t walSequence) -> std::optional<Document> {
        // Load document from WAL via persistence manager
        return persistence_->loadDocument(collection, id, 0);  // Search from beginning
    });

    // Create encryption manager
    EncryptionManager::Config encryptConfig;
    encryptConfig.enabled = config_.encryptionEnabled;
    encryptConfig.keyFilePath = config_.keyFilePath;
    encryptConfig.autoGenerateKey = config_.autoGenerateKey;
    encryption_ = std::make_unique<EncryptionManager>(encryptConfig);

    // Create event manager
    EventManager::Config eventConfig;
    eventConfig.nodeId = config_.nodeId;
    events_ = std::make_unique<EventManager>(eventConfig);

    // Create file manager
    FileManager::Config fileConfig;
    fileConfig.filesDir = config_.dataDirectory / "files";
    fileConfig.maxFileSizeMb = config_.maxFileSizeMb;
    fileConfig.allowedMimeTypes = config_.allowedFileTypes;
    fileConfig.cleanupIntervalSec = config_.fileCleanupIntervalSec;
    files_ = std::make_unique<FileManager>(fileConfig);

    // Create replication manager
    ReplicationManager::Config replConfig;
    replConfig.nodeId = config_.nodeId;
    replConfig.peerAddresses = config_.peerAddresses;
    replConfig.conflictResolution = config_.conflictResolution;
    replication_ = std::make_unique<ReplicationManager>(replConfig);

    // Wire replication entry apply callback
    replication_->setEntryApplyCallback([this](const databasepb::ReplicationEntry& entry) {
        applyReplicatedEntry(entry);
    });

    // Create gRPC implementations
    storageImpl_ = std::make_unique<DatabaseGrpcImpl>(
        *store_, *persistence_, *events_, *files_, *encryption_
    );
    replicationImpl_ = std::make_unique<DatabaseReplicationGrpcImpl>(
        *store_, *replication_, *persistence_
    );
}

void DatabaseService::applyReplicatedEntry(const databasepb::ReplicationEntry& entry) {
    try {
        switch (entry.op()) {
            case databasepb::OP_INSERT:
            case databasepb::OP_UPDATE:
            case databasepb::OP_UPSERT: {
                if (entry.data().empty()) {
                    spdlog::warn("Replication entry has no data for op {}", static_cast<int>(entry.op()));
                    return;
                }

                auto json = nlohmann::json::parse(entry.data());
                Document doc = Document::fromJson(json);
                doc.id = entry.document_id();
                doc.nodeId = entry.node_id();

                // Use loadDocument to bypass normal callbacks (avoid re-replication)
                store_->loadDocument(entry.collection(), doc);
                spdlog::trace("Applied replicated {} to {}/{} from {}",
                             entry.op() == databasepb::OP_INSERT ? "insert" :
                             entry.op() == databasepb::OP_UPDATE ? "update" : "upsert",
                             entry.collection(), entry.document_id(), entry.node_id());
                break;
            }

            case databasepb::OP_DELETE: {
                store_->remove(entry.collection(), entry.document_id());
                spdlog::trace("Applied replicated delete to {}/{} from {}",
                             entry.collection(), entry.document_id(), entry.node_id());
                break;
            }

            case databasepb::OP_CREATE_COLLECTION: {
                CollectionOptions options;
                store_->createCollection(entry.collection(), options);
                spdlog::trace("Applied replicated create collection {} from {}",
                             entry.collection(), entry.node_id());
                break;
            }

            case databasepb::OP_DROP_COLLECTION: {
                store_->dropCollection(entry.collection());
                spdlog::trace("Applied replicated drop collection {} from {}",
                             entry.collection(), entry.node_id());
                break;
            }

            default:
                spdlog::warn("Unknown replication operation type: {}", static_cast<int>(entry.op()));
                break;
        }
    } catch (const std::exception& e) {
        spdlog::error("Failed to apply replicated entry: {}", e.what());
    }
}

void DatabaseService::startGrpcServer() {
    std::string serverAddress = config_.bindAddress + ":" + std::to_string(config_.rpcPort);

    grpc::ServerBuilder builder;
    builder.AddListeningPort(serverAddress, grpc::InsecureServerCredentials());
    builder.RegisterService(storageImpl_.get());
    builder.RegisterService(replicationImpl_.get());

    // Configure server
    builder.SetMaxReceiveMessageSize(100 * 1024 * 1024);  // 100MB for file uploads
    builder.SetMaxSendMessageSize(100 * 1024 * 1024);

    grpcServer_ = builder.BuildAndStart();

    if (!grpcServer_) {
        throw std::runtime_error("Failed to start gRPC server on " + serverAddress);
    }

    spdlog::info("gRPC server listening on {}", serverAddress);
}

void DatabaseService::stopGrpcServer() {
    if (grpcServer_) {
        grpcServer_->Shutdown();
        grpcServer_.reset();
    }
}

} // namespace smartbotic::database