#include "persistent_lru_block_store.h"

namespace mediadb {

PersistentLruBlockStore::PersistentLruBlockStore(const std::string& base_dir, size_t max_bytes)
    : inner_(std::make_shared<paritypp::file_block_store>(base_dir)),
      max_bytes_(max_bytes) {
    // Re-derive byte accounting from whatever is already on disk.
    for (uint64_t gid : inner_->list_groups()) {
        for (uint32_t bidx : inner_->list_blocks(gid)) {
            std::vector<uint8_t> tmp;
            if (inner_->fetch(gid, bidx, tmp)) {
                block_sizes_[gid][bidx] = tmp.size();
                used_bytes_ += tmp.size();
            }
        }
        lru_.push_front(gid);
        lru_map_[gid] = lru_.begin();
    }

    if (max_bytes_ > 0) evict_lru();
}

size_t PersistentLruBlockStore::group_bytes(uint64_t group_id) const {
    auto it = block_sizes_.find(group_id);
    if (it == block_sizes_.end()) return 0;
    size_t total = 0;
    for (const auto& kv : it->second) total += kv.second;
    return total;
}

void PersistentLruBlockStore::touch(uint64_t group_id) {
    auto it = lru_map_.find(group_id);
    if (it != lru_map_.end()) {
        lru_.splice(lru_.begin(), lru_, it->second);
    } else {
        lru_.push_front(group_id);
        lru_map_[group_id] = lru_.begin();
    }
}

void PersistentLruBlockStore::evict_lru() {
    while (!lru_.empty() && max_bytes_ > 0 && used_bytes_ > max_bytes_) {
        uint64_t victim = lru_.back();
        used_bytes_ -= group_bytes(victim);
        block_sizes_.erase(victim);
        inner_->remove_group(victim);
        lru_map_.erase(victim);
        lru_.pop_back();
    }
}

bool PersistentLruBlockStore::store(uint64_t group_id, uint32_t block_index,
                                     const uint8_t* data, size_t len) {
    std::lock_guard<std::mutex> lock(mutex_);

    auto git = block_sizes_.find(group_id);
    if (git != block_sizes_.end()) {
        auto bit = git->second.find(block_index);
        if (bit != git->second.end())
            used_bytes_ -= bit->second;
    }

    if (!inner_->store(group_id, block_index, data, len))
        return false;

    block_sizes_[group_id][block_index] = len;
    used_bytes_ += len;
    touch(group_id);

    if (max_bytes_ > 0)
        evict_lru();
    return true;
}

bool PersistentLruBlockStore::fetch(uint64_t group_id, uint32_t block_index,
                                     std::vector<uint8_t>& out) {
    std::lock_guard<std::mutex> lock(mutex_);
    if (!inner_->fetch(group_id, block_index, out)) return false;
    touch(group_id);
    return true;
}

bool PersistentLruBlockStore::remove_group(uint64_t group_id) {
    std::lock_guard<std::mutex> lock(mutex_);
    used_bytes_ -= group_bytes(group_id);
    block_sizes_.erase(group_id);
    auto it = lru_map_.find(group_id);
    if (it != lru_map_.end()) {
        lru_.erase(it->second);
        lru_map_.erase(it);
    }
    return inner_->remove_group(group_id);
}

std::vector<uint32_t> PersistentLruBlockStore::list_blocks(uint64_t group_id) {
    return inner_->list_blocks(group_id);
}

uint32_t PersistentLruBlockStore::total_blocks() {
    return inner_->total_blocks();
}

uint32_t PersistentLruBlockStore::total_groups() {
    return inner_->total_groups();
}

std::vector<uint64_t> PersistentLruBlockStore::list_groups() {
    return inner_->list_groups();
}

bool PersistentLruBlockStore::vacuum() {
    return inner_->vacuum();
}

} // namespace mediadb
