ble-reticulum/migration/protocol_core/BLEPeerSessionManager.h

#pragma once

#include <array>
#include <cstddef>
#include <cstdint>
#include <optional>
#include <string>
#include <vector>

namespace ble_reticulum {

using PeerIdentity = std::array<uint8_t, 16>;

enum class LocalRole : uint8_t {
    Unknown,
    Central,
    Peripheral,
};

enum class InputDecision : uint8_t {
    PassToReassembler,
    ConsumedDuplicateSameIdentity,
    ConsumedDuplicateMismatchedIdentity,
    AcceptedNewIdentity,
    RejectedDuplicateIdentity,
    ErrorConsumed,
};

enum class SessionActionType : uint8_t {
    ConsumeInput,
    PassToReassembler,
    AcceptNewIdentity,
    RejectDuplicateIdentity,
    DisconnectCurrentPeer,
    DisconnectOldPeer,
    CreateFragmentationState,
    MarkPeerReady,
    UpdatePeerAddress,
    RemovePendingIdentity,
    MarkRealData,
    CleanupOldAddress,
    Warn,
};

struct ConnectionId {
    std::string address;
    uint16_t handle = 0xffff;

    bool operator==(const ConnectionId& other) const {
        return address == other.address && handle == other.handle;
    }
};

struct ConnectionSnapshot {
    ConnectionId current;
    LocalRole local_role = LocalRole::Unknown;
    std::optional<PeerIdentity> known_identity_for_address;
    std::optional<uint16_t> negotiated_mtu;

    std::optional<std::string> existing_address_for_identity;
    bool identity_has_pending_detach = false;
    bool existing_address_connected = false;
    bool existing_address_in_peer_table = false;
    bool existing_connection_is_zombie = false;
    double existing_last_real_data = 0.0;
};

struct SessionAction {
    SessionActionType type = SessionActionType::ConsumeInput;
    ConnectionId target;
    std::string old_address;
    std::string new_address;
    std::string message;
};

struct HandshakeResult {
    InputDecision decision = InputDecision::PassToReassembler;
    std::vector<SessionAction> actions;

    bool consumed = false;
    bool accepted = false;
    bool should_disconnect_current = false;
    bool should_disconnect_old = false;

    std::optional<PeerIdentity> peer_identity;
    std::string identity_key;
    std::string fragmenter_key;
    uint16_t mtu = 23;
};

struct PeerSessionView {
    PeerIdentity identity{};
    std::string identity_key;
    std::string current_address;
    uint16_t current_handle = 0xffff;
    uint16_t mtu = 23;
    bool has_fragmentation_state = false;
    bool peer_ready = false;
    double pending_identity_since = 0.0;
    double last_real_data = 0.0;
};

class BLEPeerSessionManager {
public:
    explicit BLEPeerSessionManager(double pending_identity_timeout = 30.0,
                                   double zombie_timeout = 45.0);

    HandshakeResult handleIdentityHandshake(const ConnectionSnapshot& connection,
                                            const uint8_t* data,
                                            size_t data_size,
                                            double now_seconds);

    void markConnected(const ConnectionSnapshot& connection, double now_seconds);
    void markDisconnected(const ConnectionId& connection, double now_seconds);
    void markMtu(const ConnectionId& connection, uint16_t mtu);
    void markPendingIdentity(const ConnectionId& connection, double now_seconds);
    void removePendingIdentity(const ConnectionId& connection);
    std::vector<ConnectionId> expiredPendingIdentities(double now_seconds) const;

    std::optional<PeerSessionView> sessionByAddress(const std::string& address) const;
    std::optional<PeerSessionView> sessionByIdentity(const PeerIdentity& identity) const;

    static bool isIdentityHandshakePayload(const uint8_t* data, size_t data_size);
    static PeerIdentity identityFromPayload(const uint8_t* data, size_t data_size);
    static std::string computeIdentityKey(const PeerIdentity& identity);
    static std::string computeFragmenterKey(const PeerIdentity& identity);

private:
    double pending_identity_timeout_;
    double zombie_timeout_;

    struct PendingIdentity {
        ConnectionId connection;
        double started_at = 0.0;
    };

    struct PeerSession {
        PeerIdentity identity{};
        std::string identity_key;
        std::string current_address;
        uint16_t current_handle = 0xffff;
        uint16_t mtu = 23;
        bool has_fragmentation_state = false;
        bool peer_ready = false;
        double pending_identity_since = 0.0;
        double last_real_data = 0.0;
    };

    std::vector<PeerSession> sessions_;
    std::vector<PendingIdentity> pending_identities_;

    PeerSession* findSessionByIdentityKey(const std::string& identity_key);
    const PeerSession* findSessionByIdentityKey(const std::string& identity_key) const;
    PeerSession* findSessionByAddress(const std::string& address);
    const PeerSession* findSessionByAddress(const std::string& address) const;

    void upsertAcceptedSession(const ConnectionSnapshot& connection,
                               const PeerIdentity& identity,
                               uint16_t mtu,
                               double now_seconds);
};

} // namespace ble_reticulum