#include "protocol.h"
#include <cstring>

namespace paritypp {

std::vector<uint8_t> protocol::encode_message(msg_type type,
                                               const uint8_t* payload,
                                               size_t payload_len) {
    std::vector<uint8_t> msg(HEADER_SIZE + payload_len);
    msg[0] = static_cast<uint8_t>(type);
    encode_u32(msg.data() + 1, static_cast<uint32_t>(payload_len));
    if (payload_len > 0 && payload) {
        std::memcpy(msg.data() + HEADER_SIZE, payload, payload_len);
    }
    return msg;
}

std::vector<uint8_t> protocol::encode_message(msg_type type,
                                               const std::vector<uint8_t>& payload) {
    return encode_message(type, payload.data(), payload.size());
}

bool protocol::decode_header(const uint8_t* data, size_t len,
                              msg_type& type, uint32_t& payload_len) {
    if (len < HEADER_SIZE) return false;
    type = static_cast<msg_type>(data[0]);
    payload_len = decode_u32(data + 1);
    return true;
}

void protocol::encode_u64(uint8_t* dst, uint64_t val) {
    dst[0] = static_cast<uint8_t>(val >> 56);
    dst[1] = static_cast<uint8_t>(val >> 48);
    dst[2] = static_cast<uint8_t>(val >> 40);
    dst[3] = static_cast<uint8_t>(val >> 32);
    dst[4] = static_cast<uint8_t>(val >> 24);
    dst[5] = static_cast<uint8_t>(val >> 16);
    dst[6] = static_cast<uint8_t>(val >> 8);
    dst[7] = static_cast<uint8_t>(val);
}

uint64_t protocol::decode_u64(const uint8_t* src) {
    return (static_cast<uint64_t>(src[0]) << 56) |
           (static_cast<uint64_t>(src[1]) << 48) |
           (static_cast<uint64_t>(src[2]) << 40) |
           (static_cast<uint64_t>(src[3]) << 32) |
           (static_cast<uint64_t>(src[4]) << 24) |
           (static_cast<uint64_t>(src[5]) << 16) |
           (static_cast<uint64_t>(src[6]) << 8)  |
           static_cast<uint64_t>(src[7]);
}

void protocol::encode_u32(uint8_t* dst, uint32_t val) {
    dst[0] = static_cast<uint8_t>(val >> 24);
    dst[1] = static_cast<uint8_t>(val >> 16);
    dst[2] = static_cast<uint8_t>(val >> 8);
    dst[3] = static_cast<uint8_t>(val);
}

uint32_t protocol::decode_u32(const uint8_t* src) {
    return (static_cast<uint32_t>(src[0]) << 24) |
           (static_cast<uint32_t>(src[1]) << 16) |
           (static_cast<uint32_t>(src[2]) << 8)  |
           static_cast<uint32_t>(src[3]);
}

std::vector<uint8_t> protocol::make_store_block(uint64_t group_id,
                                                  uint32_t block_index,
                                                  const uint8_t* data,
                                                  size_t data_len) {
    std::vector<uint8_t> payload(8 + 4 + data_len);
    encode_u64(payload.data(), group_id);
    encode_u32(payload.data() + 8, block_index);
    if (data_len > 0) {
        std::memcpy(payload.data() + 12, data, data_len);
    }
    return payload;
}

std::vector<uint8_t> protocol::make_fetch_block(uint64_t group_id,
                                                  uint32_t block_index) {
    std::vector<uint8_t> payload(12);
    encode_u64(payload.data(), group_id);
    encode_u32(payload.data() + 8, block_index);
    return payload;
}

std::vector<uint8_t> protocol::make_block_data(uint64_t group_id,
                                                 uint32_t block_index,
                                                 const uint8_t* data,
                                                 size_t data_len) {
    std::vector<uint8_t> payload(8 + 4 + data_len);
    encode_u64(payload.data(), group_id);
    encode_u32(payload.data() + 8, block_index);
    if (data_len > 0) {
        std::memcpy(payload.data() + 12, data, data_len);
    }
    return payload;
}

bool protocol::parse_store_block(const uint8_t* payload, size_t len,
                                   uint64_t& group_id, uint32_t& block_index,
                                   const uint8_t*& data, size_t& data_len) {
    if (len < 12) return false;
    group_id = decode_u64(payload);
    block_index = decode_u32(payload + 8);
    data = payload + 12;
    data_len = len - 12;
    return true;
}

bool protocol::parse_fetch_block(const uint8_t* payload, size_t len,
                                   uint64_t& group_id, uint32_t& block_index) {
    if (len < 12) return false;
    group_id = decode_u64(payload);
    block_index = decode_u32(payload + 8);
    return true;
}

bool protocol::parse_block_data(const uint8_t* payload, size_t len,
                                  uint64_t& group_id, uint32_t& block_index,
                                  const uint8_t*& data, size_t& data_len) {
    if (len < 12) return false;
    group_id = decode_u64(payload);
    block_index = decode_u32(payload + 8);
    data = payload + 12;
    data_len = len - 12;
    return true;
}

// --- auth protocol ---

std::vector<uint8_t> protocol::make_auth_request(const std::string& client_name) {
    if (client_name.size() > 255) return {};
    std::vector<uint8_t> payload(1 + client_name.size());
    payload[0] = static_cast<uint8_t>(client_name.size());
    std::memcpy(payload.data() + 1, client_name.data(), client_name.size());
    return payload;
}

std::vector<uint8_t> protocol::make_auth_response(const std::vector<uint8_t>& hmac) {
    return hmac; // 32 bytes HMAC-SHA256
}

std::vector<uint8_t> protocol::make_auth_challenge(const std::vector<uint8_t>& nonce) {
    return nonce; // 32 bytes nonce
}

bool protocol::parse_auth_request(const uint8_t* payload, size_t len,
                                    std::string& client_name) {
    if (len < 1) return false;
    uint8_t name_len = payload[0];
    if (len < 1 + name_len) return false;
    client_name.assign(reinterpret_cast<const char*>(payload + 1), name_len);
    return true;
}

bool protocol::parse_auth_response(const uint8_t* payload, size_t len,
                                     std::vector<uint8_t>& hmac) {
    if (len < 32) return false;
    hmac.assign(payload, payload + 32);
    return true;
}

bool protocol::parse_auth_challenge(const uint8_t* payload, size_t len,
                                      std::vector<uint8_t>& nonce) {
    if (len < 32) return false;
    nonce.assign(payload, payload + 32);
    return true;
}

// --- status protocol ---

std::vector<uint8_t> protocol::make_status_resp(const status_info& info) {
    // [1B status][4B block_count][4B group_count][4B peers_total][4B peers_online]
    std::vector<uint8_t> payload(17);
    payload[0] = static_cast<uint8_t>(info.status);
    encode_u32(payload.data() + 1, info.block_count);
    encode_u32(payload.data() + 5, info.group_count);
    encode_u32(payload.data() + 9, info.peers_total);
    encode_u32(payload.data() + 13, info.peers_online);
    return payload;
}

bool protocol::parse_status_resp(const uint8_t* payload, size_t len,
                                   status_info& info) {
    if (len < 17) return false;
    info.status = static_cast<node_status>(payload[0]);
    info.block_count = decode_u32(payload + 1);
    info.group_count = decode_u32(payload + 5);
    info.peers_total = decode_u32(payload + 9);
    info.peers_online = decode_u32(payload + 13);
    return true;
}

// --- group list protocol ---

std::vector<uint8_t> protocol::make_group_list(const std::vector<uint64_t>& group_ids) {
    // [4B count][8B group_id...]
    std::vector<uint8_t> payload(4 + group_ids.size() * 8);
    encode_u32(payload.data(), static_cast<uint32_t>(group_ids.size()));
    for (size_t i = 0; i < group_ids.size(); i++) {
        encode_u64(payload.data() + 4 + i * 8, group_ids[i]);
    }
    return payload;
}

bool protocol::parse_group_list(const uint8_t* payload, size_t len,
                                  std::vector<uint64_t>& group_ids) {
    if (len < 4) return false;
    uint32_t count = decode_u32(payload);
    if (len < 4 + count * 8) return false;
    group_ids.clear();
    group_ids.reserve(count);
    for (uint32_t i = 0; i < count; i++) {
        group_ids.push_back(decode_u64(payload + 4 + i * 8));
    }
    return true;
}

} // namespace paritypp
