#include "hmac.h"
#include "sha1.h"
#include "sha256.h"
#include "sha512.h"
#include <cstring>

template <typename HashType>
std::vector<uint8_t> hmac_impl(const std::vector<uint8_t>& key, const std::vector<uint8_t>& message) {
    std::vector<uint8_t> actual_key = key;

    if (actual_key.size() > HashType::BLOCK_SIZE) {
        HashType hash;
        hash.update(actual_key.data(), actual_key.size());
        actual_key.resize(HashType::HASH_SIZE);
        hash.finalize(actual_key.data());
    }

    if (actual_key.size() < HashType::BLOCK_SIZE) {
        actual_key.resize(HashType::BLOCK_SIZE, 0x00);
    }

    std::vector<uint8_t> o_key_pad(HashType::BLOCK_SIZE);
    std::vector<uint8_t> i_key_pad(HashType::BLOCK_SIZE);

    for (size_t i = 0; i < HashType::BLOCK_SIZE; i++) {
        o_key_pad[i] = actual_key[i] ^ 0x5c;
        i_key_pad[i] = actual_key[i] ^ 0x36;
    }

    HashType inner_hash;
    inner_hash.update(i_key_pad.data(), i_key_pad.size());
    inner_hash.update(message.data(), message.size());
    std::vector<uint8_t> inner_result(HashType::HASH_SIZE);
    inner_hash.finalize(inner_result.data());

    HashType outer_hash;
    outer_hash.update(o_key_pad.data(), o_key_pad.size());
    outer_hash.update(inner_result.data(), inner_result.size());
    std::vector<uint8_t> final_result(HashType::HASH_SIZE);
    outer_hash.finalize(final_result.data());

    return final_result;
}

std::vector<uint8_t> hmac(HashAlgorithm algo, const std::vector<uint8_t>& key, const std::vector<uint8_t>& message) {
    switch (algo) {
        case HashAlgorithm::SHA1:
            return hmac_impl<SHA1>(key, message);
        case HashAlgorithm::SHA256:
            return hmac_impl<SHA256>(key, message);
        case HashAlgorithm::SHA512:
            return hmac_impl<SHA512>(key, message);
        default:
            return std::vector<uint8_t>();
    }
}