#include "crypto/crypto.h"

#ifdef VOICECAT_HAS_NET

#include <cstring>
#include <fstream>
#include <sstream>
#include <stdexcept>

#include <mbedtls/error.h>
#include <mbedtls/sha256.h>
#include <mbedtls/ssl.h>
#include <mbedtls/x509_crt.h>

namespace voicecat::crypto {

// ── Helpers ───────────────────────────────────────────────────────────────────

static void throw_if(int rc, const char* msg) {
    if (rc != 0) {
        char buf[256];
        mbedtls_strerror(rc, buf, sizeof(buf));
        throw std::runtime_error(std::string(msg) + ": " + buf);
    }
}

static std::string compute_hex_fingerprint(const uint8_t* data, size_t len) {
    uint8_t hash[32];
    mbedtls_sha256(data, len, hash, 0);
    std::string s;
    s.reserve(64);
    const char* hex = "0123456789abcdef";
    for (auto b : hash) {
        s += hex[b >> 4];
        s += hex[b & 0xf];
    }
    return s;
}

// ── ServerIdentity ────────────────────────────────────────────────────────────

ServerIdentity ServerIdentity::generate() {
    ServerIdentity id;
    crypto_sign_ed25519_keypair(id.pk.data(), id.sk.data());
    // Fingerprint = SHA-256 of the public key
    mbedtls_sha256(id.pk.data(), id.pk.size(), id.fingerprint.data(), 0);
    return id;
}

ServerIdentity ServerIdentity::load(const std::filesystem::path& path) {
    std::ifstream f(path, std::ios::binary);
    if (!f) throw std::runtime_error("Cannot open identity file: " + path.string());
    ServerIdentity id;
    f.read(reinterpret_cast<char*>(id.pk.data()), id.pk.size());
    f.read(reinterpret_cast<char*>(id.sk.data()), id.sk.size());
    if (!f) throw std::runtime_error("Identity file truncated: " + path.string());
    mbedtls_sha256(id.pk.data(), id.pk.size(), id.fingerprint.data(), 0);
    return id;
}

void ServerIdentity::save(const std::filesystem::path& path) const {
    std::ofstream f(path, std::ios::binary | std::ios::trunc);
    if (!f) throw std::runtime_error("Cannot write identity file: " + path.string());
    f.write(reinterpret_cast<const char*>(pk.data()), pk.size());
    f.write(reinterpret_cast<const char*>(sk.data()), sk.size());
}

std::string ServerIdentity::fingerprint_hex() const {
    std::string s;
    s.reserve(96);
    const char* hex = "0123456789ABCDEF";
    for (size_t i = 0; i < fingerprint.size(); ++i) {
        if (i > 0) s += ':';
        s += hex[fingerprint[i] >> 4];
        s += hex[fingerprint[i] & 0xf];
    }
    return s;
}

// ── ServerCert ────────────────────────────────────────────────────────────────

ServerCert ServerCert::generate(const std::string& server_name) {
    mbedtls_entropy_context   entropy;
    mbedtls_ctr_drbg_context  ctr_drbg;
    mbedtls_pk_context        key;
    mbedtls_x509write_cert    cert;

    mbedtls_entropy_init(&entropy);
    mbedtls_ctr_drbg_init(&ctr_drbg);
    mbedtls_pk_init(&key);
    mbedtls_x509write_crt_init(&cert);

    try {
        const char* pers = "voicecat_cert_gen";
        throw_if(mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                                        reinterpret_cast<const unsigned char*>(pers),
                                        strlen(pers)),
                 "ctr_drbg_seed");

        // Generate ECDSA-P256 key
        throw_if(mbedtls_pk_setup(&key, mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY)),
                 "pk_setup");
        throw_if(mbedtls_ecp_gen_key(MBEDTLS_ECP_DP_SECP256R1, mbedtls_pk_ec(key),
                                      mbedtls_ctr_drbg_random, &ctr_drbg),
                 "ecp_gen_key");

        // Build self-signed cert
        mbedtls_x509write_crt_set_version(&cert, MBEDTLS_X509_CRT_VERSION_3);
        mbedtls_x509write_crt_set_md_alg(&cert, MBEDTLS_MD_SHA256);
        mbedtls_x509write_crt_set_subject_key(&cert, &key);
        mbedtls_x509write_crt_set_issuer_key(&cert, &key);

        std::string dn = "CN=" + (server_name.empty() ? std::string("voicecat") : server_name);
        throw_if(mbedtls_x509write_crt_set_subject_name(&cert, dn.c_str()), "set_subject");
        throw_if(mbedtls_x509write_crt_set_issuer_name(&cert, dn.c_str()), "set_issuer");

        // Serial = 0x01 (1 byte, value 1)
        uint8_t serial_raw[] = {0x01};
        throw_if(mbedtls_x509write_crt_set_serial_raw(&cert, serial_raw, sizeof(serial_raw)),
                 "set_serial");

        // Valid for 10 years
        throw_if(mbedtls_x509write_crt_set_validity(&cert, "20240101000000",
                                                      "20340101000000"),
                 "set_validity");
        throw_if(mbedtls_x509write_crt_set_basic_constraints(&cert, 0, -1),
                 "set_basic_constraints");

        // Write PEM cert
        unsigned char cert_buf[4096] = {};
        throw_if(mbedtls_x509write_crt_pem(&cert, cert_buf, sizeof(cert_buf),
                                             mbedtls_ctr_drbg_random, &ctr_drbg),
                 "write_cert_pem");

        // Write PEM key
        unsigned char key_buf[4096] = {};
        throw_if(mbedtls_pk_write_key_pem(&key, key_buf, sizeof(key_buf)), "write_key_pem");

        ServerCert result;
        result.pem_cert = reinterpret_cast<char*>(cert_buf);
        result.pem_key  = reinterpret_cast<char*>(key_buf);

        mbedtls_x509write_crt_free(&cert);
        mbedtls_pk_free(&key);
        mbedtls_ctr_drbg_free(&ctr_drbg);
        mbedtls_entropy_free(&entropy);
        return result;
    } catch (...) {
        mbedtls_x509write_crt_free(&cert);
        mbedtls_pk_free(&key);
        mbedtls_ctr_drbg_free(&ctr_drbg);
        mbedtls_entropy_free(&entropy);
        throw;
    }
}

ServerCert ServerCert::load(const std::filesystem::path& cert_path,
                             const std::filesystem::path& key_path) {
    auto read_file = [](const std::filesystem::path& p) -> std::string {
        std::ifstream f(p);
        if (!f) throw std::runtime_error("Cannot open: " + p.string());
        return {std::istreambuf_iterator<char>(f), {}};
    };
    ServerCert c;
    c.pem_cert = read_file(cert_path);
    c.pem_key  = read_file(key_path);
    return c;
}

void ServerCert::save(const std::filesystem::path& cert_path,
                      const std::filesystem::path& key_path) const {
    auto write_file = [](const std::filesystem::path& p, const std::string& s) {
        std::ofstream f(p, std::ios::trunc);
        if (!f) throw std::runtime_error("Cannot write: " + p.string());
        f << s;
    };
    write_file(cert_path, pem_cert);
    write_file(key_path, pem_key);
}

// ── TlsContext ────────────────────────────────────────────────────────────────

TlsContext::TlsContext(Role role, const ServerCert* server_cert,
                       const std::array<uint8_t, 32>* pinned_fp)
    : role_(role), pinned_fp_(pinned_fp) {
    mbedtls_entropy_init(&entropy_);
    mbedtls_ctr_drbg_init(&ctr_drbg_);
    mbedtls_ssl_init(&ssl_);
    mbedtls_ssl_config_init(&conf_);
    mbedtls_x509_crt_init(&srvcert_);
    mbedtls_pk_init(&pkey_);

    const char* pers = (role == Role::Server) ? "vc_server_tls" : "vc_client_tls";
    throw_if(mbedtls_ctr_drbg_seed(&ctr_drbg_, mbedtls_entropy_func, &entropy_,
                                    reinterpret_cast<const unsigned char*>(pers),
                                    strlen(pers)),
             "ctr_drbg_seed");

    int endpoint = (role == Role::Server) ? MBEDTLS_SSL_IS_SERVER : MBEDTLS_SSL_IS_CLIENT;
    throw_if(mbedtls_ssl_config_defaults(&conf_, endpoint, MBEDTLS_SSL_TRANSPORT_STREAM,
                                          MBEDTLS_SSL_PRESET_DEFAULT),
             "ssl_config_defaults");

    // TLS 1.3 only
    mbedtls_ssl_conf_min_tls_version(&conf_, MBEDTLS_SSL_VERSION_TLS1_3);
    mbedtls_ssl_conf_max_tls_version(&conf_, MBEDTLS_SSL_VERSION_TLS1_3);

    mbedtls_ssl_conf_rng(&conf_, mbedtls_ctr_drbg_random, &ctr_drbg_);

    if (role == Role::Server && server_cert) {
        // Parse server cert + key
        throw_if(mbedtls_x509_crt_parse(
                     &srvcert_,
                     reinterpret_cast<const unsigned char*>(server_cert->pem_cert.c_str()),
                     server_cert->pem_cert.size() + 1),
                 "x509_crt_parse");
        throw_if(mbedtls_pk_parse_key(
                     &pkey_,
                     reinterpret_cast<const unsigned char*>(server_cert->pem_key.c_str()),
                     server_cert->pem_key.size() + 1,
                     nullptr, 0, mbedtls_ctr_drbg_random, &ctr_drbg_),
                 "pk_parse_key");
        throw_if(mbedtls_ssl_conf_own_cert(&conf_, &srvcert_, &pkey_), "conf_own_cert");
    }

    if (role == Role::Client) {
        // Skip CA chain verification — we use TOFU via the server identity fingerprint.
        mbedtls_ssl_conf_authmode(&conf_, MBEDTLS_SSL_VERIFY_NONE);
    }

    throw_if(mbedtls_ssl_setup(&ssl_, &conf_), "ssl_setup");
}

TlsContext::~TlsContext() {
    mbedtls_ssl_close_notify(&ssl_);
    mbedtls_pk_free(&pkey_);
    mbedtls_x509_crt_free(&srvcert_);
    mbedtls_ssl_free(&ssl_);
    mbedtls_ssl_config_free(&conf_);
    mbedtls_ctr_drbg_free(&ctr_drbg_);
    mbedtls_entropy_free(&entropy_);
}

void TlsContext::set_read_timeout(uint32_t ms) {
    mbedtls_ssl_conf_read_timeout(&conf_, ms);
}

bool TlsContext::is_timeout_error(int rc) {
    return rc == MBEDTLS_ERR_SSL_TIMEOUT;
}

bool TlsContext::handshake(int socket_fd, std::string& error) {
    net_ctx_.fd = socket_fd;
    // Use the timeout-capable recv callback so set_read_timeout() takes effect.
    mbedtls_ssl_set_bio(&ssl_, &net_ctx_, mbedtls_net_send, mbedtls_net_recv,
                        mbedtls_net_recv_timeout);

    int rc;
    while ((rc = mbedtls_ssl_handshake(&ssl_)) != 0) {
        if (rc != MBEDTLS_ERR_SSL_WANT_READ && rc != MBEDTLS_ERR_SSL_WANT_WRITE) {
            char buf[256];
            mbedtls_strerror(rc, buf, sizeof(buf));
            error = buf;
            return false;
        }
    }
    ready_ = true;
    return true;
}

int TlsContext::read(uint8_t* buf, size_t len) {
    return mbedtls_ssl_read(&ssl_, buf, len);
}

int TlsContext::write(const uint8_t* buf, size_t len) {
    return mbedtls_ssl_write(&ssl_, buf, len);
}

bool TlsContext::export_keying_material(const char* label, const uint8_t* ctx, size_t ctx_len,
                                         uint8_t* out, size_t out_len) {
    return mbedtls_ssl_export_keying_material(
               &ssl_, out, out_len, label, strlen(label),
               ctx, ctx_len, ctx != nullptr) == 0;
}

bool TlsContext::peer_cert_fingerprint(std::array<uint8_t, 32>& out) const {
    if (!ready_) return false;
    const mbedtls_x509_crt* peer = mbedtls_ssl_get_peer_cert(&ssl_);
    if (!peer) return false;
    mbedtls_sha256(peer->raw.p, peer->raw.len, out.data(), 0);
    return true;
}

// ── SodiumMediaCrypto ─────────────────────────────────────────────────────────

SodiumMediaCrypto::SodiumMediaCrypto(
    const uint8_t key[crypto_aead_chacha20poly1305_ietf_KEYBYTES]) {
    std::memcpy(key_.data(), key, key_.size());
}

std::unique_ptr<SodiumMediaCrypto> SodiumMediaCrypto::derive(TlsContext& tls, uint8_t ctx_byte) {
    uint8_t key[crypto_aead_chacha20poly1305_ietf_KEYBYTES]{};
    if (!tls.export_keying_material("voicecat media v1", &ctx_byte, 1, key, sizeof(key)))
        return nullptr;
    auto p = std::make_unique<SodiumMediaCrypto>(key);
    sodium_memzero(key, sizeof(key));
    return p;
}

std::unique_ptr<SodiumMediaCrypto> SodiumMediaCrypto::derive_send(TlsContext& tls,
                                                                    bool is_client) {
    return derive(tls, is_client ? 0x00 : 0x01);
}

std::unique_ptr<SodiumMediaCrypto> SodiumMediaCrypto::derive_recv(TlsContext& tls,
                                                                    bool is_client) {
    return derive(tls, is_client ? 0x01 : 0x00);
}

void SodiumMediaCrypto::build_nonce(uint64_t counter, uint8_t nonce[12]) const {
    // nonce[0..3] = 0x00  (reserved / zero-padded)
    // nonce[4..11] = counter (big-endian u64)
    nonce[0] = nonce[1] = nonce[2] = nonce[3] = 0;
    nonce[4]  = static_cast<uint8_t>(counter >> 56);
    nonce[5]  = static_cast<uint8_t>(counter >> 48);
    nonce[6]  = static_cast<uint8_t>(counter >> 40);
    nonce[7]  = static_cast<uint8_t>(counter >> 32);
    nonce[8]  = static_cast<uint8_t>(counter >> 24);
    nonce[9]  = static_cast<uint8_t>(counter >> 16);
    nonce[10] = static_cast<uint8_t>(counter >> 8);
    nonce[11] = static_cast<uint8_t>(counter & 0xFF);
}

long SodiumMediaCrypto::seal(const uint8_t* plain, size_t len, const uint8_t* aad,
                              size_t aad_len, uint8_t* out, size_t out_cap) {
    if (out_cap < len + crypto_aead_chacha20poly1305_ietf_ABYTES) return -1;

    uint8_t nonce[crypto_aead_chacha20poly1305_ietf_NPUBBYTES];
    build_nonce(send_counter_++, nonce);

    unsigned long long sealed_len = 0;
    if (crypto_aead_chacha20poly1305_ietf_encrypt(
            out, &sealed_len, plain, static_cast<unsigned long long>(len),
            aad, static_cast<unsigned long long>(aad_len),
            nullptr, nonce, key_.data()) != 0)
        return -1;

    return static_cast<long>(sealed_len);
}

long SodiumMediaCrypto::open(const uint8_t* sealed, size_t len, const uint8_t* aad,
                              size_t aad_len, uint8_t* out, size_t out_cap) {
    if (len < crypto_aead_chacha20poly1305_ietf_ABYTES) return -1;
    if (out_cap < len - crypto_aead_chacha20poly1305_ietf_ABYTES) return -1;

    // Read the full 64-bit nonce counter directly from aad[8..15] (seq, big-endian
    // u64). Protocol v2 carries the full counter on the wire, so the nonce is exact —
    // no reconstruction/rollover guessing needed.
    if (aad_len < 16) return -1;
    uint64_t counter = (static_cast<uint64_t>(aad[8])  << 56) |
                       (static_cast<uint64_t>(aad[9])  << 48) |
                       (static_cast<uint64_t>(aad[10]) << 40) |
                       (static_cast<uint64_t>(aad[11]) << 32) |
                       (static_cast<uint64_t>(aad[12]) << 24) |
                       (static_cast<uint64_t>(aad[13]) << 16) |
                       (static_cast<uint64_t>(aad[14]) << 8)  |
                        static_cast<uint64_t>(aad[15]);

    // ── Anti-replay: REJECT-ONLY checks (no state mutation) ───────────────────
    // The counter comes from the UNAUTHENTICATED header, so we must NOT advance the
    // window before the AEAD tag is verified — otherwise a single corrupted/forged
    // packet would shove recv_highest_ far ahead and reject every later legitimate
    // packet as "too old", permanently wedging the stream. Order per RFC 3711 §3.3:
    // replay-check → authenticate → update.
    if (recv_initialized_ && counter <= recv_highest_) {
        uint64_t offset = recv_highest_ - counter;
        if (offset >= 64) return -1;                            // too old
        if (recv_window_ & (UINT64_C(1) << offset)) return -1;  // replay
    }

    uint8_t nonce[crypto_aead_chacha20poly1305_ietf_NPUBBYTES];
    build_nonce(counter, nonce);

    unsigned long long plain_len = 0;
    if (crypto_aead_chacha20poly1305_ietf_decrypt(
            out, &plain_len, nullptr, sealed, static_cast<unsigned long long>(len),
            aad, static_cast<unsigned long long>(aad_len),
            nonce, key_.data()) != 0)
        return -1;  // auth failure — leave the replay window untouched

    // ── Authenticated: now it's safe to advance the window ────────────────────
    if (!recv_initialized_) {
        recv_highest_     = counter;
        recv_window_      = 1;  // bit0 = highest itself
        recv_initialized_ = true;
    } else if (counter > recv_highest_) {
        uint64_t shift = counter - recv_highest_;
        recv_window_  = (shift >= 64) ? 0 : (recv_window_ << shift);
        recv_window_ |= 1;  // bit0 = the new highest
        recv_highest_ = counter;
    } else {
        recv_window_ |= (UINT64_C(1) << (recv_highest_ - counter));
    }

    return static_cast<long>(plain_len);
}

}  // namespace voicecat::crypto

#endif  // VOICECAT_HAS_NET
