#pragma once
#include <vector>
#include <string>
#include <functional>
#include <cstdint>

namespace netplus {

class HKDF_SHA256 {
public:
    // Callback types:
    // sha256(data) -> 32 bytes
    using Sha256Func = std::function<std::vector<uint8_t>(const std::vector<uint8_t>&)>;

    // hmac(key,data) -> 32 bytes
    using HmacFunc   = std::function<std::vector<uint8_t>(const std::vector<uint8_t>&,
                                                          const std::vector<uint8_t>&)>;

    HKDF_SHA256(Sha256Func sha256, HmacFunc hmac);

    // HKDF-Extract(salt, IKM) -> PRK
    std::vector<uint8_t> extract(const std::vector<uint8_t>& salt,
                                 const std::vector<uint8_t>& ikm) const;

    // HKDF-Expand(PRK, info, L) -> OKM
    std::vector<uint8_t> expand(const std::vector<uint8_t>& prk,
                                const std::vector<uint8_t>& info,
                                size_t outLen) const;

    // TLS 1.3 HKDF-Expand-Label
    // HKDF-Expand-Label(secret, label, context, outLen)
    std::vector<uint8_t> expand_label(const std::vector<uint8_t>& secret,
                                      const std::string& label,
                                      const std::vector<uint8_t>& context,
                                      size_t outLen) const;

    // Utility: SHA256 over empty string (used in TLS 1.3 "derived")
    std::vector<uint8_t> hash_empty() const;

private:
    Sha256Func _sha256;
    HmacFunc   _hmac;
};

// HKDF with SHA-384 for TLS_AES_256_GCM_SHA384
class HKDF_SHA384 {
public:
    // Callback types:
    // sha384(data) -> 48 bytes
    using Sha384Func = std::function<std::vector<uint8_t>(const std::vector<uint8_t>&)>;

    // hmac(key,data) -> 48 bytes
    using HmacFunc   = std::function<std::vector<uint8_t>(const std::vector<uint8_t>&,
                                                          const std::vector<uint8_t>&)>;

    HKDF_SHA384(Sha384Func sha384, HmacFunc hmac);

    // HKDF-Extract(salt, IKM) -> PRK (48 bytes)
    std::vector<uint8_t> extract(const std::vector<uint8_t>& salt,
                                 const std::vector<uint8_t>& ikm) const;

    // HKDF-Expand(PRK, info, L) -> OKM
    std::vector<uint8_t> expand(const std::vector<uint8_t>& prk,
                                const std::vector<uint8_t>& info,
                                size_t outLen) const;

    // TLS 1.3 HKDF-Expand-Label
    // HKDF-Expand-Label(secret, label, context, outLen)
    std::vector<uint8_t> expand_label(const std::vector<uint8_t>& secret,
                                      const std::string& label,
                                      const std::vector<uint8_t>& context,
                                      size_t outLen) const;

    // Utility: SHA384 over empty string (used in TLS 1.3 "derived")
    std::vector<uint8_t> hash_empty() const;

private:
    Sha384Func _sha384;
    HmacFunc   _hmac;
};

} // namespace netplus
