#include "preview_cache.h"

#include <cstring>
#include <iostream>

namespace mediadb {

// ---------- helpers ----------

uint64_t PreviewCache::cache_group_id(const std::string& key) {
    std::vector<uint8_t> input(key.begin(), key.end());
    auto hash = netplus::sha256_hash(input);
    uint64_t id = 0;
    for (int i = 0; i < 8; i++)
        id = (id << 8) | hash[i];
    return id;
}

std::vector<uint8_t> PreviewCache::encode_meta(const FragMeta& m) {
    // [4B num_chunks][8B total_size][content_type\0][etag\0]
    std::vector<uint8_t> buf;
    buf.resize(4 + 8 + m.content_type.size() + 1 + m.etag.size() + 1);
    uint8_t* p = buf.data();
    uint32_t nc = m.num_chunks;
    std::memcpy(p, &nc, 4); p += 4;
    uint64_t ts = m.total_size;
    std::memcpy(p, &ts, 8); p += 8;
    std::memcpy(p, m.content_type.data(), m.content_type.size()); p += m.content_type.size();
    *p++ = '\0';
    std::memcpy(p, m.etag.data(), m.etag.size()); p += m.etag.size();
    *p++ = '\0';
    return buf;
}

bool PreviewCache::decode_meta(const std::vector<uint8_t>& data, FragMeta& out) {
    if (data.size() < 12 + 2) return false; // min: 4+8+'\0'+'\0'
    const uint8_t* p = data.data();
    uint32_t nc;
    std::memcpy(&nc, p, 4); p += 4;
    uint64_t ts;
    std::memcpy(&ts, p, 8); p += 8;
    out.num_chunks = nc;
    out.total_size = ts;
    const char* str = reinterpret_cast<const char*>(p);
    const char* end = reinterpret_cast<const char*>(data.data() + data.size());
    // content_type
    const char* z = static_cast<const char*>(std::memchr(str, '\0', end - str));
    if (!z) return false;
    out.content_type.assign(str, z);
    str = z + 1;
    // etag
    z = static_cast<const char*>(std::memchr(str, '\0', end - str));
    if (!z) return false;
    out.etag.assign(str, z);
    return true;
}

// ---------- construction ----------

PreviewCache::PreviewCache(Cluster& cluster, size_t max_bytes_per_node) {
    store_ = std::make_shared<paritypp::memory_block_store>(max_bytes_per_node);

    // Register as store_id=1 on the existing cluster parity server
    cluster.addStore(STORE_ID, store_);

    const auto& cfg = cluster.getConfig();
    const auto& peer_map = cluster.getPeerToClientIndex();

    // Build node list reusing the cluster's peer addresses + port
    std::vector<paritypp::client::node_info> nodes;
    for (size_t pi = 0; pi < cfg.peers.size(); ++pi) {
        nodes.push_back({cfg.peers[pi].address, cfg.peers[pi].port});
    }

    paritypp::client::credentials pcreds;
    if (!cfg.client_name.empty()) {
        pcreds.client_name = cfg.client_name;
        pcreds.key.assign(cfg.client_key.begin(), cfg.client_key.end());
    }

    if (!nodes.empty()) {
        client_ = std::make_unique<paritypp::client>(
            cfg.data_blocks, cfg.parity_blocks, nodes, pcreds);
        client_->set_store_id(STORE_ID);
        if (cluster.getSelfNodeIndex() >= 0) {
            auto it = peer_map.find(static_cast<size_t>(cluster.getSelfNodeIndex()));
            if (it != peer_map.end())
                client_->set_local_node(it->second, store_);
        }
    }
}

// ---------- lifecycle ----------

void PreviewCache::start() {
    if (client_) client_->warmup();
}

// ---------- get (L2 → L1) ----------

std::shared_ptr<const BlobValue> PreviewCache::get(const std::string& key, BlobCache& l1) {
    if (!client_) return nullptr;

    // First try: single-blob retrieve
    uint64_t gid = cache_group_id(key);
    try {
        auto data = client_->retrieve(gid);
        if (!data.empty()) {
            // Check if this is fragment metadata by trying to decode
            FragMeta meta;
            if (decode_meta(data, meta) && meta.num_chunks > 0 &&
                meta.total_size > PREVIEW_FRAGMENT_SIZE) {
                // This is a meta group — reassemble from fragments
                std::vector<uint8_t> assembled;
                assembled.reserve(meta.total_size);
                for (uint32_t ci = 0; ci < meta.num_chunks; ++ci) {
                    std::string chunk_key = key + ":c:" + std::to_string(ci);
                    uint64_t cgid = cache_group_id(chunk_key);
                    auto chunk = client_->retrieve(cgid);
                    if (chunk.empty()) return nullptr; // fragment evicted, must re-render
                    assembled.insert(assembled.end(), chunk.begin(), chunk.end());

                    // Insert chunk into L1
                    BlobValue chunk_blob;
                    chunk_blob.data = std::move(chunk);
                    chunk_blob.content_type = meta.content_type;
                    chunk_blob.etag = meta.etag;
                    chunk_blob.type = BlobType::preview;
                    l1.put(chunk_key, std::move(chunk_blob));
                }

                // Insert assembled blob into L1 (if it fits, LRU handles rejection)
                BlobValue blob;
                blob.data = std::move(assembled);
                blob.content_type = meta.content_type;
                blob.etag = meta.etag;
                blob.type = BlobType::preview;
                l1.put(key, std::move(blob));
                return l1.get(key);
            }

            // Small blob — single entry. But we need content_type/etag.
            // Retrieve the meta group for metadata
            std::string meta_key = key + ":meta";
            uint64_t mgid = cache_group_id(meta_key);
            FragMeta smeta;
            try {
                auto mdata = client_->retrieve(mgid);
                if (!mdata.empty()) decode_meta(mdata, smeta);
            } catch (...) {}

            BlobValue blob;
            blob.data = std::move(data);
            blob.content_type = smeta.content_type;
            blob.etag = smeta.etag;
            blob.type = BlobType::preview;
            l1.put(key, std::move(blob));
            return l1.get(key);
        }
    } catch (...) {
        // L2 unreachable — fall through to miss
    }

    return nullptr;
}

// ---------- put (render result → L1 + L2) ----------

void PreviewCache::put(const std::string& key, const BlobValue& blob, BlobCache& l1) {
    // Always store metadata
    FragMeta meta;
    meta.total_size = blob.data.size();
    meta.content_type = blob.content_type;
    meta.etag = blob.etag;

    if (blob.data.size() <= PREVIEW_FRAGMENT_SIZE) {
        // Small preview — single group
        meta.num_chunks = 0;
        l1.put(key, BlobValue{blob.data, blob.content_type, blob.etag, BlobType::preview});

        if (client_) {
            try {
                // Store metadata
                std::string meta_key = key + ":meta";
                uint64_t mgid = cache_group_id(meta_key);
                auto mdata = encode_meta(meta);
                client_->store(mgid, mdata);

                // Store data
                uint64_t gid = cache_group_id(key);
                client_->store(gid, blob.data);
            } catch (const std::exception& e) {
                std::cerr << "[PCACHE] L2 store failed: " << e.what() << "\n";
            }
        }
    } else {
        // Large preview — fragment into 10 MB chunks
        size_t total = blob.data.size();
        uint32_t num_chunks = static_cast<uint32_t>(
            (total + PREVIEW_FRAGMENT_SIZE - 1) / PREVIEW_FRAGMENT_SIZE);
        meta.num_chunks = num_chunks;

        for (uint32_t ci = 0; ci < num_chunks; ++ci) {
            size_t offset = static_cast<size_t>(ci) * PREVIEW_FRAGMENT_SIZE;
            size_t len = std::min(PREVIEW_FRAGMENT_SIZE, total - offset);

            std::string chunk_key = key + ":c:" + std::to_string(ci);

            // L1: each fragment as separate cache entry
            BlobValue chunk_blob;
            chunk_blob.data.assign(blob.data.begin() + offset,
                                    blob.data.begin() + offset + len);
            chunk_blob.content_type = blob.content_type;
            chunk_blob.etag = blob.etag;
            chunk_blob.type = BlobType::preview;
            l1.put(chunk_key, std::move(chunk_blob));

            // L2: each fragment as separate parity group
            if (client_) {
                try {
                    uint64_t cgid = cache_group_id(chunk_key);
                    client_->store(cgid, blob.data.data() + offset, len);
                } catch (const std::exception& e) {
                    std::cerr << "[PCACHE] L2 chunk store failed: " << e.what() << "\n";
                }
            }
        }

        // Also store the assembled blob in L1 (LRU will reject if too large)
        l1.put(key, BlobValue{blob.data, blob.content_type, blob.etag, BlobType::preview});

        // Store fragment metadata in L2
        if (client_) {
            try {
                uint64_t gid = cache_group_id(key);
                auto mdata = encode_meta(meta);
                client_->store(gid, mdata);
            } catch (const std::exception& e) {
                std::cerr << "[PCACHE] L2 meta store failed: " << e.what() << "\n";
            }
        }
    }
}

} // namespace mediadb
