//
// Created by Kirill Zhukov on 20.04.2025.
//

#include "UDB.h"

namespace usub::core
{
    UDB::UDB(const std::string& db_name, const std::string& shm_name,
             size_t shm_queue_capacity, size_t max_memtable_size, bool create_new)
        : db_name_(db_name),
          shm_name_(shm_name),
          shm_queue_capacity_(shm_queue_capacity),
          max_memtable_size_(max_memtable_size),
          shm_manager_(shm_name, sizeof(SharedCommandQueue), create_new),
          version_manager_(db_name),
          memtable_manager_(db_name + "_wal", max_memtable_size, this->version_manager_),
          compactor_(this->version_manager_),
          running_(true)
    {
        ::shm_unlink(("/" + shm_name).c_str());
        new(this->shm_manager_.base_ptr()) SharedCommandQueue(this->shm_queue_capacity_);
        this->command_queue_ = reinterpret_cast<SharedCommandQueue*>(this->shm_manager_.base_ptr());

        recover_from_logs();

        this->background_flush_thread_ = std::thread(&UDB::background_flush_worker, this);
        compactor_.run();
    }

    UDB::~UDB()
    {
        running_ = false;

        if (this->background_flush_thread_.joinable())
            this->background_flush_thread_.join();

    }

    void UDB::recover_from_logs()
    {
        utils::RecoveryLog recovery_log(db_name_);

        recovery_log.replay([this](const utils::Hash128& key, const std::string& value, bool is_tombstone)
        {
            if (!is_tombstone)
            {
                this->fast_cache_[key] = value;
                this->memtable_manager_.put(key, value);
            }
            else
            {
                this->fast_cache_.erase(key);
                this->memtable_manager_.remove(key);
            }
        });

        std::cout << "[UDB] Recovery complete for database: " << this->db_name_ << "\n";
    }

    void UDB::run()
    {
        std::cout << "[UDB] Server for DB '" << this->db_name_ << "' started.\n";

        while (this->running_)
        {
            auto opt_cmd = this->command_queue_->try_pop();

            if (!opt_cmd)
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(1));
                continue;
            }

            process_command(*opt_cmd);
        }
    }

    void UDB::process_command(const Command& cmd)
    {
        switch (cmd.op)
        {
        case OperationType::PUT:
            {
                std::string value(cmd.value, cmd.value_size);
                this->fast_cache_[cmd.key] = value;
                this->memtable_manager_.put(cmd.key, value);
                break;
            }
        case OperationType::DELETE:
            {
                fast_cache_.erase(cmd.key);
                memtable_manager_.remove(cmd.key);
                break;
            }
        case OperationType::FIND:
            {
                auto it = fast_cache_.find(cmd.key);
                if (it != fast_cache_.end())
                    std::cout << "[FIND] " << usub::utils::to_string(cmd.key) << " => " << it->second << "\n";
                else
                    std::cout << "[FIND] " << usub::utils::to_string(cmd.key) << " not found\n";
                break;
            }
        default:
            std::cout << "[UNKNOWN COMMAND]\n";
            break;
        }
    }

    void UDB::background_flush_worker()
    {
        while (running_)
        {
            memtable_manager_.flush_batch();
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
    }
}