#include "crypto.h"

#include <algorithm>
#include <filesystem>
#include <fstream>
#include <sstream>
#include <zstd.h>

#include <sodium.h>

#include "util.h"

namespace m2pack
{

namespace
{

std::uint8_t from_hex_digit(const char ch)
{
    if (ch >= '0' && ch <= '9')
    {
        return static_cast<std::uint8_t>(ch - '0');
    }
    if (ch >= 'a' && ch <= 'f')
    {
        return static_cast<std::uint8_t>(10 + ch - 'a');
    }
    if (ch >= 'A' && ch <= 'F')
    {
        return static_cast<std::uint8_t>(10 + ch - 'A');
    }
    fail("Invalid hex digit");
}

}  // namespace

void crypto_init()
{
    if (sodium_init() < 0)
    {
        fail("libsodium initialization failed");
    }
}

std::array<std::uint8_t, kHashSize> hash_bytes(const std::vector<std::uint8_t>& bytes)
{
    std::array<std::uint8_t, kHashSize> out {};
    crypto_generichash(out.data(), out.size(), bytes.data(), bytes.size(), nullptr, 0);
    return out;
}

std::array<std::uint8_t, kHashSize> hash_string(std::string_view text)
{
    std::array<std::uint8_t, kHashSize> out {};
    crypto_generichash(out.data(), out.size(),
                       reinterpret_cast<const unsigned char*>(text.data()), text.size(),
                       nullptr, 0);
    return out;
}

std::vector<std::uint8_t> random_bytes(std::size_t size)
{
    std::vector<std::uint8_t> out(size);
    randombytes_buf(out.data(), out.size());
    return out;
}

std::array<std::uint8_t, kAeadNonceSize> random_nonce()
{
    std::array<std::uint8_t, kAeadNonceSize> nonce {};
    randombytes_buf(nonce.data(), nonce.size());
    return nonce;
}

std::vector<std::uint8_t> load_hex_file(const std::string& path)
{
    std::ifstream input(path);
    if (!input)
    {
        fail("Failed to open key file: " + path);
    }

    std::stringstream buffer;
    buffer << input.rdbuf();
    std::string text = buffer.str();

    text.erase(std::remove_if(text.begin(), text.end(), [](unsigned char ch) {
        return std::isspace(ch) != 0;
    }), text.end());

    if (text.size() % 2 != 0)
    {
        fail("Hex file has odd length: " + path);
    }

    std::vector<std::uint8_t> bytes(text.size() / 2);
    for (std::size_t i = 0; i < bytes.size(); ++i)
    {
        bytes[i] = static_cast<std::uint8_t>((from_hex_digit(text[i * 2]) << 4) | from_hex_digit(text[i * 2 + 1]));
    }

    return bytes;
}

std::string to_hex(const std::uint8_t* data, std::size_t size)
{
    static constexpr char kHex[] = "0123456789abcdef";

    std::string out;
    out.resize(size * 2);
    for (std::size_t i = 0; i < size; ++i)
    {
        out[i * 2] = kHex[(data[i] >> 4) & 0x0f];
        out[i * 2 + 1] = kHex[data[i] & 0x0f];
    }
    return out;
}

std::string to_hex(const std::vector<std::uint8_t>& data)
{
    return to_hex(data.data(), data.size());
}

std::vector<std::uint8_t> compress_zstd(const std::vector<std::uint8_t>& input)
{
    const auto bound = ZSTD_compressBound(input.size());
    std::vector<std::uint8_t> out(bound);
    const auto written = ZSTD_compress(out.data(), out.size(), input.data(), input.size(), 12);
    if (ZSTD_isError(written))
    {
        fail("ZSTD compression failed");
    }
    out.resize(written);
    return out;
}

std::vector<std::uint8_t> decompress_zstd(const std::vector<std::uint8_t>& input, std::size_t output_size)
{
    std::vector<std::uint8_t> out(output_size);
    const auto written = ZSTD_decompress(out.data(), out.size(), input.data(), input.size());
    if (ZSTD_isError(written))
    {
        fail("ZSTD decompression failed");
    }
    if (written != output_size)
    {
        fail("ZSTD decompression size mismatch");
    }
    return out;
}

std::vector<std::uint8_t> encrypt_payload(
    const std::vector<std::uint8_t>& plaintext,
    const std::array<std::uint8_t, kAeadKeySize>& key,
    const std::array<std::uint8_t, kAeadNonceSize>& nonce,
    std::string_view associated_data)
{
    std::vector<std::uint8_t> ciphertext(
        plaintext.size() + crypto_aead_xchacha20poly1305_ietf_ABYTES);

    unsigned long long written = 0;
    if (crypto_aead_xchacha20poly1305_ietf_encrypt(
            ciphertext.data(),
            &written,
            plaintext.data(),
            plaintext.size(),
            reinterpret_cast<const unsigned char*>(associated_data.data()),
            associated_data.size(),
            nullptr,
            nonce.data(),
            key.data()) != 0)
    {
        fail("Encryption failed");
    }

    ciphertext.resize(static_cast<std::size_t>(written));
    return ciphertext;
}

std::vector<std::uint8_t> decrypt_payload(
    const std::vector<std::uint8_t>& ciphertext,
    const std::array<std::uint8_t, kAeadKeySize>& key,
    const std::array<std::uint8_t, kAeadNonceSize>& nonce,
    std::string_view associated_data)
{
    std::vector<std::uint8_t> plaintext(ciphertext.size());

    unsigned long long written = 0;
    if (crypto_aead_xchacha20poly1305_ietf_decrypt(
            plaintext.data(),
            &written,
            nullptr,
            ciphertext.data(),
            ciphertext.size(),
            reinterpret_cast<const unsigned char*>(associated_data.data()),
            associated_data.size(),
            nonce.data(),
            key.data()) != 0)
    {
        fail("Payload authentication failed");
    }

    plaintext.resize(static_cast<std::size_t>(written));
    return plaintext;
}

std::vector<std::uint8_t> encrypt_payload_stream(
    const std::vector<std::uint8_t>& plaintext,
    const std::array<std::uint8_t, kAeadKeySize>& key,
    const std::array<std::uint8_t, kAeadNonceSize>& nonce)
{
    std::vector<std::uint8_t> ciphertext(plaintext.size());
    if (!plaintext.empty())
    {
        crypto_stream_xchacha20_xor(
            ciphertext.data(),
            plaintext.data(),
            plaintext.size(),
            nonce.data(),
            key.data());
    }
    return ciphertext;
}

std::vector<std::uint8_t> decrypt_payload_stream(
    const std::vector<std::uint8_t>& ciphertext,
    const std::array<std::uint8_t, kAeadKeySize>& key,
    const std::array<std::uint8_t, kAeadNonceSize>& nonce)
{
    std::vector<std::uint8_t> plaintext(ciphertext.size());
    if (!ciphertext.empty())
    {
        crypto_stream_xchacha20_xor(
            plaintext.data(),
            ciphertext.data(),
            ciphertext.size(),
            nonce.data(),
            key.data());
    }
    return plaintext;
}

std::vector<std::uint8_t> sign_detached(
    const std::vector<std::uint8_t>& data,
    const std::vector<std::uint8_t>& secret_key)
{
    if (secret_key.size() != crypto_sign_SECRETKEYBYTES)
    {
        fail("Signing secret key has invalid size");
    }

    std::vector<std::uint8_t> signature(crypto_sign_BYTES);
    unsigned long long sig_len = 0;
    crypto_sign_detached(signature.data(), &sig_len, data.data(), data.size(), secret_key.data());
    signature.resize(static_cast<std::size_t>(sig_len));
    return signature;
}

bool verify_detached(
    const std::vector<std::uint8_t>& data,
    const std::uint8_t* signature,
    const std::vector<std::uint8_t>& public_key)
{
    if (public_key.size() != crypto_sign_PUBLICKEYBYTES)
    {
        fail("Signing public key has invalid size");
    }

    return crypto_sign_verify_detached(signature, data.data(), data.size(), public_key.data()) == 0;
}

}  // namespace m2pack