#include "preview_cache.h"
#include "persistent_lru_block_store.h"

#include <cstring>
#include <iostream>

namespace mediadb {

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

namespace {
// Tags every envelope so get() can tell a Small (metadata+payload) record
// apart from a Fragmented (metadata-only) record, and reject anything else
// retrieved from the same group slot outright. Without this, binary video
// containers (mp4/webm) routinely "decode" as a plausible-looking oversized
// FragMeta, since their leading bytes parse as a huge num_chunks/total_size
// — causing get() to attempt a bogus fragment reassembly and miss every time.
constexpr uint32_t ENVELOPE_MAGIC = 0x31435046; // "PFC1"
} // namespace

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_frag_meta(const FragMeta& m) {
    // [4B magic][1B kind][4B num_chunks][8B total_size][content_type\0][etag\0]
    std::vector<uint8_t> buf;
    buf.resize(4 + 1 + 4 + 8 + m.content_type.size() + 1 + m.etag.size() + 1);
    uint8_t* p = buf.data();
    uint32_t magic = ENVELOPE_MAGIC;
    std::memcpy(p, &magic, 4); p += 4;
    *p++ = static_cast<uint8_t>(EnvelopeKind::Fragmented);
    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_frag_meta(const std::vector<uint8_t>& data, FragMeta& out) {
    if (data.size() < 4 + 1 + 12 + 2) return false; // magic+kind+4+8+'\0'+'\0'
    const uint8_t* p = data.data();
    uint32_t magic;
    std::memcpy(&magic, p, 4); p += 4;
    if (magic != ENVELOPE_MAGIC) return false;
    if (*p++ != static_cast<uint8_t>(EnvelopeKind::Fragmented)) return false;
    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;
}

std::vector<uint8_t> PreviewCache::encode_small(const std::string& content_type,
                                                 const std::string& etag,
                                                 const std::vector<uint8_t>& data) {
    // [4B magic][1B kind][8B data_len][content_type\0][etag\0][data...]
    std::vector<uint8_t> buf;
    buf.resize(4 + 1 + 8 + content_type.size() + 1 + etag.size() + 1 + data.size());
    uint8_t* p = buf.data();
    uint32_t magic = ENVELOPE_MAGIC;
    std::memcpy(p, &magic, 4); p += 4;
    *p++ = static_cast<uint8_t>(EnvelopeKind::Small);
    uint64_t dlen = data.size();
    std::memcpy(p, &dlen, 8); p += 8;
    std::memcpy(p, content_type.data(), content_type.size()); p += content_type.size();
    *p++ = '\0';
    std::memcpy(p, etag.data(), etag.size()); p += etag.size();
    *p++ = '\0';
    if (!data.empty()) std::memcpy(p, data.data(), data.size());
    return buf;
}

bool PreviewCache::decode_small(const std::vector<uint8_t>& data, std::string& content_type,
                                 std::string& etag, std::vector<uint8_t>& payload) {
    if (data.size() < 4 + 1 + 8 + 2) return false; // magic+kind+8+'\0'+'\0'
    const uint8_t* p = data.data();
    uint32_t magic;
    std::memcpy(&magic, p, 4); p += 4;
    if (magic != ENVELOPE_MAGIC) return false;
    if (*p++ != static_cast<uint8_t>(EnvelopeKind::Small)) return false;
    uint64_t dlen;
    std::memcpy(&dlen, p, 8); p += 8;
    const char* str = reinterpret_cast<const char*>(p);
    const char* end = reinterpret_cast<const char*>(data.data() + data.size());
    const char* z = static_cast<const char*>(std::memchr(str, '\0', end - str));
    if (!z) return false;
    content_type.assign(str, z);
    str = z + 1;
    z = static_cast<const char*>(std::memchr(str, '\0', end - str));
    if (!z) return false;
    etag.assign(str, z);
    const uint8_t* payload_start = reinterpret_cast<const uint8_t*>(z + 1);
    size_t payload_avail = static_cast<size_t>(end - reinterpret_cast<const char*>(payload_start));
    if (dlen != payload_avail) return false; // truncated/corrupt record
    payload.assign(payload_start, payload_start + dlen);
    return true;
}

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

PreviewCache::PreviewCache(Cluster& cluster, size_t max_bytes_per_node,
                            const std::string& store_path) {
    if (!store_path.empty()) {
        store_ = std::make_shared<PersistentLruBlockStore>(store_path, max_bytes_per_node);
    } else {
        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;

    // Single retrieve — small previews carry their metadata inline, so a
    // hit here is a full hit with no second round trip.
    uint64_t gid = cache_group_id(key);
    try {
        auto data = client_->retrieve(gid);
        if (data.empty()) return nullptr;

        FragMeta meta;
        if (decode_frag_meta(data, meta)) {
            // Fragmented — reassemble from chunk groups
            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

                BlobValue chunk_blob;
                chunk_blob.content_type = meta.content_type;
                chunk_blob.etag = meta.etag;
                chunk_blob.type = BlobType::preview;
                assembled.insert(assembled.end(), chunk.begin(), chunk.end());
                chunk_blob.data = std::move(chunk);
                l1.put(chunk_key, std::move(chunk_blob));
            }

            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);
        }

        std::string content_type, etag;
        std::vector<uint8_t> payload;
        if (!decode_small(data, content_type, etag, payload)) return nullptr; // unrecognized envelope

        BlobValue blob;
        blob.data = std::move(payload);
        blob.content_type = std::move(content_type);
        blob.etag = std::move(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) {
    if (blob.data.size() <= PREVIEW_FRAGMENT_SIZE) {
        // Small preview — metadata + payload in one group, one round trip
        l1.put(key, BlobValue{blob.data, blob.content_type, blob.etag, BlobType::preview});

        if (client_) {
            try {
                uint64_t gid = cache_group_id(key);
                auto envelope = encode_small(blob.content_type, blob.etag, blob.data);
                client_->store(gid, envelope);
            } catch (const std::exception& e) {
                std::cerr << "[PCACHE] L2 store failed: " << e.what() << "\n";
            }
        }
    } else {
        // Large preview — fragment into 10 MB chunks
        FragMeta meta;
        meta.total_size = blob.data.size();
        meta.content_type = blob.content_type;
        meta.etag = blob.etag;

        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_frag_meta(meta);
                client_->store(gid, mdata);
            } catch (const std::exception& e) {
                std::cerr << "[PCACHE] L2 meta store failed: " << e.what() << "\n";
            }
        }
    }
}

} // namespace mediadb
