#pragma once

// ==============================================================================
// Auth + crypto primitives
// ==============================================================================
// Port target: medusalix/xone auth/crypto.c + auth/crypto.h.
//
// The kernel driver uses the in-kernel crypto API (crypto_shash for
// SHA-256/HMAC, crypto_akcipher for PKCS#1 RSA, crypto_kpp for ECDH
// P-256). These are replaced here with user-space equivalents:
//   - SHA-256 / HMAC-SHA256: CommonCrypto (part of libSystem)
//   - RSA PKCS#1 v1.5:       Security.framework
//   - ECDH P-256:            self-contained implementation (no public
//                            macOS C API exists for EC key agreement)
//   - randomness:            arc4random_buf
//
// AES-CCMP frame encryption itself runs on the MT76 chip; this layer
// only derives keys and installs them (see xone_mt76_set_client_key).
// ==============================================================================

#include <array>
#include <cstdint>
#include <span>
#include <string_view>

#include <CommonCrypto/CommonDigest.h>

namespace xone::auth {

using u8 = std::uint8_t;

inline constexpr auto k_sha256_len = 32;

// --------------------------------------------------------------------------
// SHA-256
// --------------------------------------------------------------------------
// Incremental SHA-256. Instances are cheap to copy; a copy is an
// independent snapshot of the current hash state, which the handshake
// uses to compute the running transcript without losing progress.
class sha256 {
public:
    sha256();
    ~sha256() = default;
    sha256(sha256 const&) = default;
    auto operator=(sha256 const&) -> sha256& = default;

    auto update(std::span<u8 const> data) -> void;
    auto finalize(std::array<u8, k_sha256_len>& out) -> void;

private:
    CC_SHA256_CTX ctx_;
};

// --------------------------------------------------------------------------
// HMAC-SHA256
// --------------------------------------------------------------------------
class hmac_sha256 {
public:
    explicit hmac_sha256(std::span<u8 const> key);

    auto update(std::span<u8 const> data) -> void;
    auto finalize(std::array<u8, k_sha256_len>& out) -> void;

private:
    sha256 inner_;
    sha256 outer_;
};

// --------------------------------------------------------------------------
// TLS P_SHA256 PRF
// --------------------------------------------------------------------------
// Expands key + label + seed into out_len bytes, exactly like the kernel
// driver's gip_auth_compute_prf(). Used to derive the master secret, the
// handshake transcript checks, and the session key.
auto prf_sha256(std::span<u8 const> key, std::string_view label,
                std::span<u8 const> seed, std::span<u8> out) -> void;

// --------------------------------------------------------------------------
// RSA (PKCS#1 v1.5)
// --------------------------------------------------------------------------
// Encrypts plaintext with a DER RSAPublicKey (the 270-byte ASN.1 SEQUENCE
// that Microsoft controllers embed in their X.509 certificate). out must
// be at least the RSA modulus size (256 bytes for 2048-bit keys).
auto rsa_encrypt_pkcs1(std::span<u8 const> der_key,
                       std::span<u8 const> plaintext, std::span<u8> out)
    -> bool;

// --------------------------------------------------------------------------
// ECDH (P-256)
// --------------------------------------------------------------------------
using ec_scalar = std::array<u8, 32>;  // big-endian private key / x-coord
using ec_point = std::array<u8, 64>;   // X || Y, big-endian, no 0x04 prefix

// q = d * G.
auto ec_base_point_multiply(ec_scalar const& d, ec_point& q) -> void;

// shared = x-coordinate of d * peer (the standard P-256 shared secret).
// Returns false if peer is not a valid point on the curve.
auto ec_compute_shared(ec_scalar const& d, ec_point const& peer,
                       ec_scalar& shared) -> bool;

// Generate a random keypair.
auto ec_generate_keypair(ec_scalar& d, ec_point& q) -> void;

// --------------------------------------------------------------------------
// Randomness
// --------------------------------------------------------------------------
auto random_bytes(std::span<u8> out) -> void;

} // namespace xone::auth