/*
 * crypto/crypto.h — TLS 1.3 (mbedTLS) and the media AEAD (libsodium).
 *
 * Design: docs/security.md. Control channel = TLS 1.3. Media = keys exported from the TLS
 * session (RFC 5705 / 8446) + per-frame ChaCha20-Poly1305 with a counter nonce and a
 * sliding-window replay filter. Encryption is MANDATORY — never add a plaintext path.
 */
#ifndef VOICECAT_CRYPTO_CRYPTO_H
#define VOICECAT_CRYPTO_CRYPTO_H

#include <cstddef>
#include <cstdint>
#include <string>

#ifdef VOICECAT_HAS_NET

#include <array>
#include <filesystem>
#include <functional>
#include <memory>

// libsodium
#include <sodium.h>

// mbedTLS
#include <mbedtls/ctr_drbg.h>
#include <mbedtls/entropy.h>
#include <mbedtls/net_sockets.h>
#include <mbedtls/pk.h>
#include <mbedtls/ssl.h>
#include <mbedtls/x509_crt.h>

namespace voicecat::crypto {

// ── Server identity ────────────────────────────────────────────────────────────
// Long-lived Ed25519 key identifying this server instance across cert rotations.
// Fingerprint is the 32-byte SHA-256 of the public key.
struct ServerIdentity {
    std::array<uint8_t, crypto_sign_ed25519_PUBLICKEYBYTES> pk{};
    std::array<uint8_t, crypto_sign_ed25519_SECRETKEYBYTES> sk{};
    std::array<uint8_t, 32> fingerprint{};

    static ServerIdentity generate();
    static ServerIdentity load(const std::filesystem::path& path);
    void save(const std::filesystem::path& path) const;
    std::string fingerprint_hex() const;
};

// ── Server TLS certificate ─────────────────────────────────────────────────────
// Self-signed ECDSA-P256 cert for TLS. On first run, generated and persisted.
struct ServerCert {
    std::string pem_cert;
    std::string pem_key;

    static ServerCert generate(const std::string& server_name);
    static ServerCert load(const std::filesystem::path& cert_path,
                           const std::filesystem::path& key_path);
    void save(const std::filesystem::path& cert_path,
              const std::filesystem::path& key_path) const;
};

// ── TLS 1.3 context ───────────────────────────────────────────────────────────
// Wraps mbedTLS for one TLS connection (server or client side).
// All public methods except close() must be called from a single thread at a time.
class TlsContext {
   public:
    enum class Role { Server, Client };

    // server_cert: required for server role; nullptr for client
    // pinned_fp: 32-byte Ed25519 fingerprint to accept (client TOFU); nullptr = any
    TlsContext(Role role, const ServerCert* server_cert,
               const std::array<uint8_t, 32>* pinned_fp = nullptr);
    ~TlsContext();

    TlsContext(const TlsContext&) = delete;
    TlsContext& operator=(const TlsContext&) = delete;

    // Perform the TLS handshake over an already-connected BSD socket fd.
    // Blocking — run from a WorkerPool thread.
    // Returns true on success; error contains a diagnostic string on failure.
    bool handshake(int socket_fd, std::string& error);

    // Read/write post-handshake (single-threaded). Returns bytes transferred, or <0 on error.
    int read(uint8_t* buf, size_t len);
    int write(const uint8_t* buf, size_t len);

    // RFC 5705 / RFC 8446 §7.5 exporter — derive media keys after handshake.
    bool export_keying_material(const char* label, const uint8_t* ctx, size_t ctx_len,
                                uint8_t* out, size_t out_len);

    // M4 TOFU: SHA-256 of the peer's leaf X.509 certificate (DER), valid only after a
    // successful Role::Client handshake(). This is the value vc_client pins — see
    // voicecat.h's vc_tofu_status doc comment for why the cert fingerprint is pinned instead
    // of the declared Ed25519 server_identity_fingerprint. Returns false if no peer cert is
    // available (e.g. Role::Server, or handshake() hasn't succeeded).
    bool peer_cert_fingerprint(std::array<uint8_t, 32>& out) const;

    // Whether the handshake completed.
    bool ready() const { return ready_; }

    // Underlying socket fd (valid after handshake). For select() in the caller.
    int native_fd() const { return net_ctx_.fd; }

    // Set per-read timeout (ms, 0 = blocking). Affects post-handshake reads.
    void set_read_timeout(uint32_t ms);

    // True when the given return value from read() indicates a read timeout.
    static bool is_timeout_error(int rc);

   private:
    Role role_;
    const std::array<uint8_t, 32>* pinned_fp_;
    bool ready_{false};

    mbedtls_entropy_context   entropy_{};
    mbedtls_ctr_drbg_context  ctr_drbg_{};
    mbedtls_ssl_context       ssl_{};
    mbedtls_ssl_config        conf_{};
    mbedtls_x509_crt          srvcert_{};
    mbedtls_pk_context        pkey_{};
    mbedtls_net_context       net_ctx_{};
};

// ── Media AEAD (M2) ───────────────────────────────────────────────────────────
// Per-frame voice encryption. Abstracted so the backend is swappable.
class MediaCrypto {
   public:
    virtual ~MediaCrypto() = default;
    // Encrypt plain[0..len) with AAD aad[0..aad_len). Write ciphertext+MAC to out.
    // out_cap must be >= len + crypto_aead_chacha20poly1305_ietf_ABYTES (16).
    // Returns total bytes written on success, or -1 on error.
    virtual long seal(const uint8_t* plain, size_t len, const uint8_t* aad, size_t aad_len,
                      uint8_t* out, size_t out_cap) = 0;
    // Decrypt+authenticate sealed[0..len). len includes the 16-byte MAC.
    // Returns number of plaintext bytes written to out, or -1 on auth failure / replay.
    virtual long open(const uint8_t* sealed, size_t len, const uint8_t* aad, size_t aad_len,
                      uint8_t* out, size_t out_cap) = 0;
};

// ── ChaCha20-Poly1305 backend ──────────────────────────────────────────────────
// Keys are derived from the TLS session via RFC 5705 / mbedTLS exporter.
// Nonce scheme: 4 zero bytes ‖ monotonic-counter(u64 big-endian, 8 bytes).
// Anti-replay: 64-bit sliding window keyed on the received counter.
class SodiumMediaCrypto final : public MediaCrypto {
   public:
    // ctx_byte: 0x00 = client→server direction, 0x01 = server→client direction.
    static std::unique_ptr<SodiumMediaCrypto> derive(TlsContext& tls, uint8_t ctx_byte);

    // Convenience: derive the key used to encrypt outgoing frames.
    // is_client=true  → ctx=0x00 (client sends); is_client=false → ctx=0x01 (server sends).
    static std::unique_ptr<SodiumMediaCrypto> derive_send(TlsContext& tls, bool is_client);
    // Convenience: derive the key used to decrypt incoming frames.
    // is_client=true  → ctx=0x01 (client recvs); is_client=false → ctx=0x00 (server recvs).
    static std::unique_ptr<SodiumMediaCrypto> derive_recv(TlsContext& tls, bool is_client);

    // Unit-test constructor: supply a raw 32-byte key directly.
    explicit SodiumMediaCrypto(const uint8_t key[crypto_aead_chacha20poly1305_ietf_KEYBYTES]);

    // Returns the send counter for the NEXT seal() call (use as frame seq).
    uint64_t peek_send_counter() const { return send_counter_; }

    // seal(): increments send_counter_; nonce derived from internal counter.
    long seal(const uint8_t* plain, size_t len, const uint8_t* aad, size_t aad_len,
              uint8_t* out, size_t out_cap) override;

    // open(): reads the full 64-bit counter from aad[8..15] (seq field), checks
    // anti-replay, then decrypts. The replay window is advanced ONLY after the AEAD
    // tag verifies, so a corrupted/forged packet cannot poison it (RFC 3711 §3.3).
    long open(const uint8_t* sealed, size_t len, const uint8_t* aad, size_t aad_len,
              uint8_t* out, size_t out_cap) override;

   private:
    void build_nonce(uint64_t counter, uint8_t nonce[12]) const;

    std::array<uint8_t, crypto_aead_chacha20poly1305_ietf_KEYBYTES> key_{};

    // Send state (used only in seal()).
    uint64_t send_counter_{0};

    // Receive anti-replay state (used only in open()).
    // Window: highest accepted counter + bitmask of last 64 accepted counters.
    uint64_t recv_highest_{0};         // highest counter seen and accepted
    uint64_t recv_window_{0};          // bit i set → (highest - i) was accepted
    bool     recv_initialized_{false}; // first packet initializes the window
};

}  // namespace voicecat::crypto

#else  // !VOICECAT_HAS_NET — skeleton stubs

namespace voicecat::crypto {

class MediaCrypto {
   public:
    virtual ~MediaCrypto() = default;
    virtual long seal(const uint8_t*, size_t, const uint8_t*, size_t, uint8_t*, size_t) = 0;
    virtual long open(const uint8_t*, size_t, const uint8_t*, size_t, uint8_t*, size_t) = 0;
};

class SodiumMediaCrypto final : public MediaCrypto {
   public:
    long seal(const uint8_t*, size_t, const uint8_t*, size_t, uint8_t*, size_t) override { return -1; }
    long open(const uint8_t*, size_t, const uint8_t*, size_t, uint8_t*, size_t) override { return -1; }
    uint64_t peek_send_counter() const { return 0; }
};

}  // namespace voicecat::crypto

#endif  // VOICECAT_HAS_NET
#endif  // VOICECAT_CRYPTO_CRYPTO_H
