#pragma once

// ==============================================================================
// GIP authentication handshake
// ==============================================================================
// Port target: medusalix/xone auth/auth.c + auth/auth.h.
//
// The kernel driver runs the RSA/ECDH exchanges on workqueues and speaks to
// the controller through gip_send_authenticate(). Here the exchanges run
// synchronously and the GIP layer is reached through the auth_sink callback
// interface, which the gip::client implements.
// ==============================================================================

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

#include "auth/crypto.hpp"

namespace xone::auth {

using u8 = std::uint8_t;
using u16 = std::uint16_t;

// Wire sizes from the upstream auth.h.
inline constexpr auto k_trailer_len = 8;
inline constexpr auto k_random_len = 32;
inline constexpr auto k_certificate_max_len = 1024;
inline constexpr auto k_pubkey_len = 270;       // v1 RSA public key
inline constexpr auto k_secret_len = 48;
inline constexpr auto k_encrypted_pms_len = 256;
inline constexpr auto k_transcript_len = 32;
inline constexpr auto k_session_key_len = 16;
inline constexpr auto k_pubkey2_len = 64;       // v2 EC public key (X || Y)
inline constexpr auto k_secret2_len = 32;

// Callback interface implemented by the GIP layer.
class auth_sink {
public:
    virtual ~auth_sink() = default;

    // Send an AUTHENTICATE command packet to the controller.
    virtual auto send(std::span<u8 const> pkt, bool acknowledge) -> int = 0;

    // Install the derived AES-CCMP session key (16 bytes).
    virtual auto set_encryption_key(std::span<u8 const> key) -> int = 0;
};

// Runs the handshake for one controller. Port of struct gip_auth.
class auth {
public:
    explicit auth(auth_sink& sink);

    // Begin the handshake by sending the (v1) host hello.
    auto start() -> int;

    // Process one incoming AUTHENTICATE packet from the controller.
    auto process_pkt(std::span<u8 const> data) -> int;

private:
    auto send_pkt(u8 cmd, void const* pkt, std::size_t len) -> int;
    auto request_pkt(u8 cmd, std::uint16_t len) -> int;
    auto send_hello() -> int;
    auto send_hello2() -> int;
    auto send_finish(u8 cmd) -> int;
    auto send_complete() -> int;
    auto exchange_rsa() -> void;
    auto exchange_ecdh() -> void;
    auto complete_handshake() -> void;
    auto handle_pkt_acknowledge() -> int;
    auto handle_pkt_data(std::span<u8 const> data) -> int;
    auto run_deferred_action() -> void;
    auto dispatch_pkt(u8 cmd, std::span<u8 const> data) -> int;
    auto handle_hello(std::span<u8 const> data) -> int;
    auto handle_certificate(std::span<u8 const> data) -> int;
    auto handle_finish(std::span<u8 const> data) -> int;
    auto handle_hello2(std::span<u8 const> data) -> int;
    auto handle_certificate2(std::span<u8 const> data) -> int;
    auto handle_pubkey(std::span<u8 const> data) -> int;
    auto get_transcript() -> std::array<u8, k_transcript_len>;

    auth_sink& m_sink;
    sha256 m_transcript;
    u8 m_last_sent_command = 0;

    std::array<u8, k_random_len> m_random_host{};
    std::array<u8, k_random_len> m_random_client{};
    std::array<u8, k_pubkey_len> m_pubkey_client{};
    std::array<u8, k_pubkey2_len> m_pubkey_client2{};
    std::array<u8, k_secret_len> m_master_secret{};

    enum class deferred_action {
        none,
        exchange_rsa,
        exchange_ecdh,
        complete,
    };
    deferred_action m_deferred_action = deferred_action::none;
};

} // namespace xone::auth