#include "hkdf.h"
#include <iostream>
#include <iomanip>
#include <stdexcept>
#include <cstring>

namespace netplus {

static inline void dumpHex(const char* prefix, const uint8_t* p, size_t n, size_t max = 64) {
    std::ios oldState(nullptr);
    oldState.copyfmt(std::cerr);

    std::cerr << prefix << " (" << n << " bytes)";
    if (n > max) std::cerr << " showing first " << max;
    std::cerr << ":\n  ";

    size_t m = (n < max) ? n : max;
    for (size_t i = 0; i < m; ++i) {
        std::cerr << std::hex << std::setw(2) << std::setfill('0') << int(p[i]) << ' ';
        if ((i % 16) == 15 && i + 1 < m) std::cerr << "\n  ";
    }
    std::cerr << std::dec << "\n";

    std::cerr.copyfmt(oldState);
}

HKDF_SHA256::HKDF_SHA256(Sha256Func sha256, HmacFunc hmac)
    : _sha256(std::move(sha256)), _hmac(std::move(hmac))
{
    if (!_sha256 || !_hmac)
        throw std::runtime_error("HKDF_SHA256: callbacks missing");
}

std::vector<uint8_t> HKDF_SHA256::extract(const std::vector<uint8_t>& salt,
                                          const std::vector<uint8_t>& ikm) const
{
    // RFC 5869: PRK = HMAC(salt, IKM)
    return _hmac(salt, ikm);
}

std::vector<uint8_t> HKDF_SHA256::expand(const std::vector<uint8_t>& prk,
                                         const std::vector<uint8_t>& info,
                                         size_t outLen) const
{
    // RFC 5869 Expand
    if (outLen == 0) return {};

    const size_t hashLen = 32;
    if (outLen > 255 * hashLen)
        throw std::runtime_error("HKDF expand: outLen too large");

    std::vector<uint8_t> okm;
    okm.reserve(outLen);

    std::vector<uint8_t> t;
    uint8_t counter = 1;

    while (okm.size() < outLen) {
        std::vector<uint8_t> in;
        in.reserve(t.size() + info.size() + 1);

        in.insert(in.end(), t.begin(), t.end());
        in.insert(in.end(), info.begin(), info.end());
        in.push_back(counter);

        t = _hmac(prk, in); // 32 bytes

        size_t need = std::min(t.size(), outLen - okm.size());
        okm.insert(okm.end(), t.begin(), t.begin() + need);

        counter++;
        if (counter == 0)
            throw std::runtime_error("HKDF expand: counter overflow");
    }

    return okm;
}

std::vector<uint8_t> HKDF_SHA256::expand_label(const std::vector<uint8_t>& secret,
                                               const std::string& label,
                                               const std::vector<uint8_t>& context,
                                               size_t outLen) const
{
    // TLS 1.3 HKDF-Expand-Label:
    //
    // struct {
    //   uint16 length = outLen;
    //   opaque label<7..255> = "tls13 " + label;
    //   opaque context<0..255> = context;
    // } HkdfLabel;
    //
    // return HKDF-Expand(secret, HkdfLabel, outLen)

    std::string fullLabel = "tls13 " + label;
    if (fullLabel.size() > 255)
        throw std::runtime_error("HKDF expand_label: label too long");
    if (context.size() > 255)
        throw std::runtime_error("HKDF expand_label: context too long");
    if (outLen > 0xFFFF)
        throw std::runtime_error("HKDF expand_label: outLen too big");

    std::vector<uint8_t> info;
    info.reserve(2 + 1 + fullLabel.size() + 1 + context.size());

    // length
    info.push_back(uint8_t((outLen >> 8) & 0xFF));
    info.push_back(uint8_t(outLen & 0xFF));

    // label
    info.push_back(uint8_t(fullLabel.size()));
    info.insert(info.end(), fullLabel.begin(), fullLabel.end());

    // context
    info.push_back(uint8_t(context.size()));
    info.insert(info.end(), context.begin(), context.end());

    return expand(secret, info, outLen);
}

std::vector<uint8_t> HKDF_SHA256::hash_empty() const
{
    std::vector<uint8_t> empty;
    return _sha256(empty);
}

// ============================================================================
// HKDF_SHA384 Implementation
// ============================================================================

HKDF_SHA384::HKDF_SHA384(Sha384Func sha384, HmacFunc hmac)
    : _sha384(std::move(sha384)), _hmac(std::move(hmac))
{
    if (!_sha384 || !_hmac)
        throw std::runtime_error("HKDF_SHA384: callbacks missing");
}

std::vector<uint8_t> HKDF_SHA384::extract(const std::vector<uint8_t>& salt,
                                          const std::vector<uint8_t>& ikm) const
{
    // RFC 5869: PRK = HMAC(salt, IKM)
    return _hmac(salt, ikm);
}

std::vector<uint8_t> HKDF_SHA384::expand(const std::vector<uint8_t>& prk,
                                         const std::vector<uint8_t>& info,
                                         size_t outLen) const
{
    // RFC 5869 Expand
    if (outLen == 0) return {};

    const size_t hashLen = 48;  // SHA-384 output is 48 bytes
    if (outLen > 255 * hashLen)
        throw std::runtime_error("HKDF_SHA384 expand: outLen too large");

    std::vector<uint8_t> okm;
    okm.reserve(outLen);

    std::vector<uint8_t> t;
    uint8_t counter = 1;

    while (okm.size() < outLen) {
        std::vector<uint8_t> in;
        in.reserve(t.size() + info.size() + 1);

        in.insert(in.end(), t.begin(), t.end());
        in.insert(in.end(), info.begin(), info.end());
        in.push_back(counter);

        t = _hmac(prk, in); // 48 bytes

        size_t need = std::min(t.size(), outLen - okm.size());
        okm.insert(okm.end(), t.begin(), t.begin() + need);

        counter++;
        if (counter == 0)
            throw std::runtime_error("HKDF_SHA384 expand: counter overflow");
    }

    return okm;
}

std::vector<uint8_t> HKDF_SHA384::expand_label(const std::vector<uint8_t>& secret,
                                               const std::string& label,
                                               const std::vector<uint8_t>& context,
                                               size_t outLen) const
{
    // TLS 1.3 HKDF-Expand-Label (same structure as SHA-256 version)
    std::string fullLabel = "tls13 " + label;
    if (fullLabel.size() > 255)
        throw std::runtime_error("HKDF_SHA384 expand_label: label too long");
    if (context.size() > 255)
        throw std::runtime_error("HKDF_SHA384 expand_label: context too long");
    if (outLen > 0xFFFF)
        throw std::runtime_error("HKDF_SHA384 expand_label: outLen too big");

    std::vector<uint8_t> info;
    info.reserve(2 + 1 + fullLabel.size() + 1 + context.size());

    // length
    info.push_back(uint8_t((outLen >> 8) & 0xFF));
    info.push_back(uint8_t(outLen & 0xFF));

    // label
    info.push_back(uint8_t(fullLabel.size()));
    info.insert(info.end(), fullLabel.begin(), fullLabel.end());

    // context
    info.push_back(uint8_t(context.size()));
    info.insert(info.end(), context.begin(), context.end());

    return expand(secret, info, outLen);
}

std::vector<uint8_t> HKDF_SHA384::hash_empty() const
{
    std::vector<uint8_t> empty;
    return _sha384(empty);
}

} // namespace netplus

