/* SPDX-License-Identifier: GPL-3.0-only */
/* Role-aware SSH stream transport with two fixed broker sessions. */

#include "ssh_transport.h"

#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <stdio.h>
#include <string.h>

#include "admin_ssh_console.h"
#include "esp_heap_caps.h"
#include "esp_log.h"
#include "esp_timer.h"
#include "freertos/FreeRTOS.h"
#include "freertos/semphr.h"
#include "freertos/task.h"
#include "lwip/inet.h"
#include "lwip/sockets.h"
#include "lwip/tcp.h"
#include "sdkconfig.h"
#include "secure_random.h"
#include "serial_service.h"
#include "ssh_security.h"
#include "user_database.h"
#include <wolfssl/wolfcrypt/memory.h>
#include <wolfssl/wolfcrypt/random.h>
#include <wolfssh/ssh.h>

#if defined(CONFIG_MBEDTLS_HARDWARE_AES) && CONFIG_MBEDTLS_HARDWARE_AES
#error "Concurrent mbedTLS/wolfSSH operation requires mbedTLS software AES"
#endif
#if !defined(NO_WOLFSSL_ESP32_CRYPT_AES) || \
    !defined(NO_WOLFSSL_ESP32_CRYPT_HASH)
#error "wolfSSH AES/SHA must not use independently locked ESP32 crypto hardware"
#endif

#define SSH_TRANSPORT_TASK_STACK_SIZE 20480U
#define SSH_TRANSPORT_TASK_PRIORITY 5U
#define SSH_TRANSPORT_TASK_CORE 1
#define SSH_TRANSPORT_LOOP_DELAY_MS 10U
#define SSH_TRANSPORT_COMMAND_TIMEOUT_MS 10000U
#define SSH_TRANSPORT_RECONCILE_INTERVAL_US 250000LL
#define SSH_TRANSPORT_LISTEN_BACKLOG 2
#define SSH_TRANSPORT_ACCEPT_BUDGET 4U
#define SSH_TRANSPORT_MAX_PENDING_HANDSHAKES SSH_TRANSPORT_MAX_SESSIONS
#define SSH_TRANSPORT_MAX_AUTH_ATTEMPTS 3U
#define SSH_TRANSPORT_GENERATION_MAX 0x3fffffffU
#define SSH_TRANSPORT_WOLFSSH_READ_BUDGET 2048U

static const char *TAG = "ssh_transport";

typedef struct {
    ssh_transport_session_state_t state;
    uint32_t generation;
    uint32_t session_id;
    int socket_fd;
    WOLFSSH *ssh;
    session_broker_client_id_t broker_client_id;
    ssh_transport_session_route_t route;
    user_principal_t principal;
    user_principal_t pending_principal;
    bool principal_valid;
    bool pending_principal_valid;
    bool authenticated;
    bool shell_requested;
    uint8_t authentication_attempts;
    word32 io_read_budget;
    bool writer;
    bool close_requested;
    int64_t handshake_deadline_us;
    int64_t last_reconcile_us;
    size_t rx_offset;
    size_t rx_length;
    size_t tx_offset;
    size_t tx_length;
    uint8_t rx_buffer[SSH_TRANSPORT_IO_BUFFER_SIZE];
    uint8_t tx_buffer[SSH_TRANSPORT_IO_BUFFER_SIZE];
    char peer[48];
} ssh_slot_t;

static portMUX_TYPE s_lock = portMUX_INITIALIZER_UNLOCKED;
static TaskHandle_t s_task;
static ssh_slot_t s_slots[SSH_TRANSPORT_MAX_SESSIONS];
static ssh_transport_session_snapshot_t
    s_session_snapshots[SSH_TRANSPORT_MAX_SESSIONS];
static uint32_t s_external_close_id[SSH_TRANSPORT_MAX_SESSIONS];
static ssh_transport_counters_t s_counters;
static SemaphoreHandle_t s_command_mutex;
static bool s_initializing;
static bool s_init_faulted;
static bool s_initialized;
static bool s_running;
static bool s_transitioning;
static bool s_desired_running;
static bool s_cleanup_pending;
static uint32_t s_requested_sequence;
static uint32_t s_completed_sequence;
static esp_err_t s_command_result = ESP_ERR_INVALID_STATE;
static esp_err_t s_last_error = ESP_ERR_INVALID_STATE;

/* Owned exclusively by the transport task. */
static WOLFSSH_CTX *s_context;
static int s_listen_fd = -1;

static void add_counter(uint64_t *counter, uint64_t value)
{
    taskENTER_CRITICAL(&s_lock);
    *counter += value;
    taskEXIT_CRITICAL(&s_lock);
}

static uint32_t next_generation(uint32_t generation)
{
    return generation >= SSH_TRANSPORT_GENERATION_MAX ? 1U : generation + 1U;
}

static uint32_t make_session_id(size_t slot_index, uint32_t generation)
{
    return (generation << 2U) | (uint32_t)(slot_index + 1U);
}

static void notify_task(void)
{
    if (s_task != NULL) {
        xTaskNotifyGive(s_task);
    }
}

static admin_ssh_console_token_t admin_console_token(const ssh_slot_t *slot,
                                                     size_t slot_index)
{
    return (admin_ssh_console_token_t){
        .slot_index = (uint8_t)slot_index,
        .session_id = slot->session_id,
        .slot_generation = slot->generation,
    };
}

static void publish_slot(const ssh_slot_t *slot, size_t slot_index)
{
    ssh_transport_session_snapshot_t snapshot = {
        .active = slot->state != SSH_TRANSPORT_SESSION_FREE,
        .authenticated = slot->authenticated,
        .principal_valid = slot->principal_valid,
        .writer = slot->writer,
        .close_requested = slot->close_requested,
        .rx_pending = slot->rx_length > slot->rx_offset,
        .tx_pending = slot->tx_length > slot->tx_offset,
        .session_id = slot->session_id,
        .generation = slot->generation,
        .socket_fd = slot->socket_fd,
        .broker_client_id = slot->broker_client_id,
        .state = slot->state,
        .route = slot->route,
        .user_role = slot->principal_valid ? slot->principal.role : USER_ROLE_USER,
        .auth_method = slot->principal_valid
                           ? slot->principal.method
                           : USER_AUTH_METHOD_PASSWORD,
    };
    if (slot->route == SSH_TRANSPORT_ROUTE_ADMIN_CONSOLE) {
        admin_ssh_console_session_snapshot_t admin_snapshot;
        admin_ssh_console_token_t token = admin_console_token(slot, slot_index);
        if (admin_ssh_console_get_session_snapshot(&token, &admin_snapshot) == ESP_OK) {
            snapshot.admin_command_pending = admin_snapshot.command_pending;
            snapshot.admin_output_pending = (uint32_t)admin_snapshot.output_length;
            snapshot.tx_pending = snapshot.tx_pending || admin_snapshot.output_pending;
        }
    }
    if (slot->principal_valid) {
        memcpy(snapshot.username, slot->principal.username,
               slot->principal.username_length);
    }
    memcpy(snapshot.peer, slot->peer, sizeof(snapshot.peer));

    taskENTER_CRITICAL(&s_lock);
    s_session_snapshots[slot_index] = snapshot;
    taskEXIT_CRITICAL(&s_lock);
}

static bool consume_external_close(const ssh_slot_t *slot, size_t slot_index)
{
    taskENTER_CRITICAL(&s_lock);
    bool requested = s_external_close_id[slot_index] != 0U &&
                     s_external_close_id[slot_index] == slot->session_id;
    if (requested || slot->state == SSH_TRANSPORT_SESSION_FREE) {
        s_external_close_id[slot_index] = 0U;
    }
    taskEXIT_CRITICAL(&s_lock);
    return requested;
}

static void *ssh_malloc(size_t size)
{
    return heap_caps_malloc_prefer(size, 2,
                                   MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT,
                                   MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT);
}

static void ssh_free(void *pointer)
{
    heap_caps_free(pointer);
}

static void *ssh_realloc(void *pointer, size_t size)
{
    return heap_caps_realloc_prefer(pointer, size, 2,
                                    MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT,
                                    MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT);
}

static int ssh_seed(OS_Seed *seed, byte *output, word32 size)
{
    (void)seed;
    return secure_random_fill(output, size) == ESP_OK ? 0 : -1;
}

static int bounded_ssh_receive(WOLFSSH *ssh, void *data, word32 size,
                               void *context)
{
    (void)ssh;
    ssh_slot_t *slot = (ssh_slot_t *)context;
    if (slot == NULL || data == NULL || slot->socket_fd < 0) {
        return WS_CBIO_ERR_GENERAL;
    }
    if (slot->io_read_budget == 0U) {
        return WS_CBIO_ERR_WANT_READ;
    }

    word32 request_size = size;
    if (request_size > slot->io_read_budget) {
        request_size = slot->io_read_budget;
    }
    int received = recv(slot->socket_fd, data, request_size, 0);
    if (received > 0) {
        slot->io_read_budget -= (word32)received;
        return received;
    }
    if (received == 0) {
        return WS_CBIO_ERR_CONN_CLOSE;
    }
    if (errno == EAGAIN || errno == EWOULDBLOCK) {
        return WS_CBIO_ERR_WANT_READ;
    }
    if (errno == EINTR) {
        return WS_CBIO_ERR_ISR;
    }
    if (errno == ECONNRESET) {
        return WS_CBIO_ERR_CONN_RST;
    }
    if (errno == ECONNABORTED) {
        return WS_CBIO_ERR_CONN_CLOSE;
    }
    return WS_CBIO_ERR_GENERAL;
}

static int allowed_auth_types(WOLFSSH *ssh, void *context)
{
    (void)ssh;
    (void)context;
    return WOLFSSH_USERAUTH_PASSWORD | WOLFSSH_USERAUTH_PUBLICKEY;
}

static void clear_pending_principal(ssh_slot_t *slot)
{
    if (slot != NULL) {
        secure_wipe(&slot->pending_principal, sizeof(slot->pending_principal));
        slot->pending_principal_valid = false;
    }
}

static bool complete_authentication_attempt(ssh_slot_t *slot, bool failed)
{
    add_counter(&s_counters.authentication_attempts, 1U);
    if (failed) {
        add_counter(&s_counters.authentication_failures, 1U);
    }
    if (slot != NULL && slot->authentication_attempts < UINT8_MAX) {
        ++slot->authentication_attempts;
    }
    if (!failed || slot == NULL ||
        slot->authentication_attempts < SSH_TRANSPORT_MAX_AUTH_ATTEMPTS) {
        return true;
    }

    slot->close_requested = true;
    if (slot->socket_fd >= 0) {
        (void)shutdown(slot->socket_fd, SHUT_RDWR);
    }
    return false;
}

static int authenticate_password(ssh_slot_t *slot,
                                 WS_UserAuthData *authentication)
{
    clear_pending_principal(slot);
    if (authentication->sf.password.hasNewPassword != 0U) {
        (void)complete_authentication_attempt(slot, true);
        return slot->close_requested ? WOLFSSH_USERAUTH_REJECTED
                                     : WOLFSSH_USERAUTH_INVALID_AUTHTYPE;
    }

    user_principal_t principal;
    bool authenticated = false;
    esp_err_t error = user_database_authenticate_password(
        authentication->username, authentication->usernameSz,
        authentication->sf.password.password,
        authentication->sf.password.passwordSz,
        &principal, &authenticated);
    if (error == ESP_OK && authenticated) {
        (void)complete_authentication_attempt(slot, false);
        slot->principal = principal;
        slot->principal_valid = true;
        slot->authenticated = true;
        return WOLFSSH_USERAUTH_SUCCESS;
    }

    secure_wipe(&principal, sizeof(principal));
    bool retry = complete_authentication_attempt(slot, true);
    if (!retry) {
        return WOLFSSH_USERAUTH_REJECTED;
    }
    return error == ESP_OK ? WOLFSSH_USERAUTH_INVALID_PASSWORD
                           : WOLFSSH_USERAUTH_FAILURE;
}

static int authenticate_public_key(ssh_slot_t *slot,
                                   WS_UserAuthData *authentication)
{
    WS_UserAuthData_PublicKey *public_key = &authentication->sf.publicKey;
    clear_pending_principal(slot);

    user_principal_t principal;
    bool authorized = false;
    esp_err_t error = ESP_OK;
    if (public_key->isCert != 0U) {
        memset(&principal, 0, sizeof(principal));
    } else {
        error = user_database_authorize_ssh_public_key(
            authentication->username, authentication->usernameSz,
            public_key->publicKeyType, public_key->publicKeyTypeSz,
            public_key->publicKey, public_key->publicKeySz,
            &principal, &authorized);
    }

    if (error != ESP_OK || !authorized) {
        secure_wipe(&principal, sizeof(principal));
        if (public_key->hasSignature == 0U) {
            return error == ESP_OK ? WOLFSSH_USERAUTH_INVALID_PUBLICKEY
                                   : WOLFSSH_USERAUTH_FAILURE;
        }
        bool retry = complete_authentication_attempt(slot, true);
        if (!retry) {
            return WOLFSSH_USERAUTH_REJECTED;
        }
        return error == ESP_OK ? WOLFSSH_USERAUTH_INVALID_PUBLICKEY
                               : WOLFSSH_USERAUTH_FAILURE;
    }

    if (public_key->hasSignature != 0U) {
        slot->pending_principal = principal;
        slot->pending_principal_valid = true;
    }
    secure_wipe(&principal, sizeof(principal));
    return WOLFSSH_USERAUTH_SUCCESS;
}

static int authenticate_user(byte authentication_type,
                             WS_UserAuthData *authentication,
                             void *context)
{
    ssh_slot_t *slot = (ssh_slot_t *)context;
    if (slot == NULL || authentication == NULL ||
        authentication_type != authentication->type) {
        clear_pending_principal(slot);
        return WOLFSSH_USERAUTH_INVALID_AUTHTYPE;
    }
    if (authentication_type == WOLFSSH_USERAUTH_PASSWORD) {
        return authenticate_password(slot, authentication);
    }
    if (authentication_type == WOLFSSH_USERAUTH_PUBLICKEY) {
        return authenticate_public_key(slot, authentication);
    }

    clear_pending_principal(slot);
    return WOLFSSH_USERAUTH_INVALID_AUTHTYPE;
}

static int authentication_result(byte result, WS_UserAuthData *authentication,
                                 void *context)
{
    ssh_slot_t *slot = (ssh_slot_t *)context;
    if (slot == NULL || authentication == NULL ||
        authentication->type != WOLFSSH_USERAUTH_PUBLICKEY ||
        authentication->sf.publicKey.hasSignature == 0U) {
        clear_pending_principal(slot);
        return WS_ERROR;
    }

    if (result != WOLFSSH_USERAUTH_SUCCESS) {
        (void)complete_authentication_attempt(slot, true);
        clear_pending_principal(slot);
        return WS_SUCCESS;
    }

    bool current = false;
    esp_err_t error = slot->pending_principal_valid
                          ? user_database_principal_is_current(
                                &slot->pending_principal, &current)
                          : ESP_ERR_INVALID_STATE;
    if (error != ESP_OK || !current) {
        (void)complete_authentication_attempt(slot, true);
        clear_pending_principal(slot);
        return WS_ERROR;
    }

    (void)complete_authentication_attempt(slot, false);
    slot->principal = slot->pending_principal;
    slot->principal_valid = true;
    slot->authenticated = true;
    clear_pending_principal(slot);
    return WS_SUCCESS;
}

static int accept_shell(WOLFSSH_CHANNEL *channel, void *context)
{
    (void)channel;
    ssh_slot_t *slot = (ssh_slot_t *)context;
    if (slot == NULL) {
        return 1;
    }
    slot->shell_requested = true;
    return 0;
}

static int reject_channel_request(WOLFSSH_CHANNEL *channel, void *context)
{
    (void)channel;
    (void)context;
    add_counter(&s_counters.request_rejections, 1U);
    return 1;
}

static bool wolfssh_would_block(WOLFSSH *ssh, int result)
{
    int error = ssh != NULL ? wolfSSH_get_error(ssh) : result;
    return result == WS_WANT_READ || result == WS_WANT_WRITE ||
           result == WS_REKEYING || result == WS_WINDOW_FULL ||
           result == WS_CHAN_RXD ||
           error == WS_WANT_READ || error == WS_WANT_WRITE ||
           error == WS_REKEYING || error == WS_WINDOW_FULL ||
           error == WS_CHAN_RXD;
}


static void close_socket(int *socket_fd)
{
    if (*socket_fd >= 0) {
        shutdown(*socket_fd, SHUT_RDWR);
        close(*socket_fd);
        *socket_fd = -1;
    }
}

static bool cleanup_slot(ssh_slot_t *slot)
{
    if (slot->route == SSH_TRANSPORT_ROUTE_ADMIN_CONSOLE) {
        size_t slot_index = (size_t)(slot - s_slots);
        admin_ssh_console_token_t token = admin_console_token(slot, slot_index);
        admin_ssh_console_close(&token);
    }
    if (slot->ssh != NULL) {
        (void)wolfSSH_shutdown(slot->ssh);
        wolfSSH_free(slot->ssh);
        slot->ssh = NULL;
    }
    close_socket(&slot->socket_fd);

    if (slot->broker_client_id != SESSION_BROKER_NO_CLIENT) {
        esp_err_t error = session_broker_disconnect(slot->broker_client_id);
        if (error != ESP_OK && error != ESP_ERR_NOT_FOUND) {
            add_counter(&s_counters.broker_failures, 1U);
            slot->state = SSH_TRANSPORT_SESSION_CLOSING;
            return false;
        }
        slot->broker_client_id = SESSION_BROKER_NO_CLIENT;
        add_counter(&s_counters.disconnections, 1U);
    }

    uint32_t generation = slot->generation;
    memset(slot, 0, sizeof(*slot));
    slot->state = SSH_TRANSPORT_SESSION_FREE;
    slot->generation = generation;
    slot->socket_fd = -1;
    return true;
}

static void request_slot_close(ssh_slot_t *slot, bool revoked)
{
    if (slot->state == SSH_TRANSPORT_SESSION_FREE) {
        return;
    }
    if (!slot->close_requested && revoked) {
        add_counter(&s_counters.session_revocations, 1U);
    }
    slot->close_requested = true;
}

static esp_err_t set_nonblocking(int socket_fd)
{
    int flags = fcntl(socket_fd, F_GETFL, 0);
    if (flags < 0 || fcntl(socket_fd, F_SETFL, flags | O_NONBLOCK) < 0) {
        return ESP_FAIL;
    }
    return ESP_OK;
}

static esp_err_t create_context(void)
{
    uint8_t private_key[SSH_SECURITY_PRIVATE_KEY_DER_CAPACITY] = {0};
    size_t private_key_length = 0U;
    esp_err_t error = ssh_security_copy_private_key(
        private_key, sizeof(private_key), &private_key_length);
    if (error != ESP_OK) {
        secure_wipe(private_key, sizeof(private_key));
        return error;
    }

    WOLFSSH_CTX *context = wolfSSH_CTX_new(WOLFSSH_ENDPOINT_SERVER, NULL);
    if (context == NULL) {
        secure_wipe(private_key, sizeof(private_key));
        return ESP_ERR_NO_MEM;
    }

    int result = wolfSSH_CTX_UsePrivateKey_buffer(
        context, private_key, (word32)private_key_length, WOLFSSH_FORMAT_ASN1);
    secure_wipe(private_key, sizeof(private_key));
    if (result != WS_SUCCESS) {
        wolfSSH_CTX_free(context);
        return ESP_FAIL;
    }

    wolfSSH_SetIORecv(context, bounded_ssh_receive);
    wolfSSH_SetUserAuth(context, authenticate_user);
    wolfSSH_SetUserAuthTypes(context, allowed_auth_types);
    wolfSSH_SetUserAuthResult(context, authentication_result);
    (void)wolfSSH_CTX_SetChannelReqShellCb(context, accept_shell);
    (void)wolfSSH_CTX_SetChannelReqExecCb(context, reject_channel_request);
    (void)wolfSSH_CTX_SetChannelReqSubsysCb(context, reject_channel_request);

    s_context = context;
    return ESP_OK;
}

static esp_err_t create_listener(void)
{
    int socket_fd = socket(AF_INET, SOCK_STREAM, IPPROTO_IP);
    if (socket_fd < 0) {
        return ESP_FAIL;
    }

    int enabled = 1;
    (void)setsockopt(socket_fd, SOL_SOCKET, SO_REUSEADDR,
                     &enabled, sizeof(enabled));
    struct sockaddr_in address = {
        .sin_family = AF_INET,
        .sin_port = htons(SSH_TRANSPORT_PORT),
        .sin_addr.s_addr = htonl(INADDR_ANY),
    };
    if (bind(socket_fd, (struct sockaddr *)&address, sizeof(address)) < 0 ||
        listen(socket_fd, SSH_TRANSPORT_LISTEN_BACKLOG) < 0 ||
        set_nonblocking(socket_fd) != ESP_OK) {
        close(socket_fd);
        return ESP_FAIL;
    }

    s_listen_fd = socket_fd;
    return ESP_OK;
}

static esp_err_t start_runtime(void)
{
    esp_err_t error = create_context();
    if (error == ESP_OK) {
        error = create_listener();
    }
    if (error != ESP_OK) {
        close_socket(&s_listen_fd);
        if (s_context != NULL) {
            wolfSSH_CTX_free(s_context);
            s_context = NULL;
        }
    }
    return error;
}

static esp_err_t stop_runtime(void)
{
    close_socket(&s_listen_fd);
    for (size_t index = 0U; index < SSH_TRANSPORT_MAX_SESSIONS; ++index) {
        request_slot_close(&s_slots[index], false);
    }

    bool all_free = false;
    for (unsigned int attempt = 0U; attempt < 100U; ++attempt) {
        all_free = true;
        for (size_t index = 0U; index < SSH_TRANSPORT_MAX_SESSIONS; ++index) {
            ssh_slot_t *slot = &s_slots[index];
            if (slot->state != SSH_TRANSPORT_SESSION_FREE) {
                all_free = false;
                (void)cleanup_slot(slot);
                publish_slot(slot, index);
            }
        }
        if (all_free) {
            break;
        }
        vTaskDelay(pdMS_TO_TICKS(10U));
    }

    all_free = true;
    for (size_t index = 0U; index < SSH_TRANSPORT_MAX_SESSIONS; ++index) {
        if (s_slots[index].state != SSH_TRANSPORT_SESSION_FREE) {
            all_free = false;
            break;
        }
    }
    if (s_context != NULL) {
        wolfSSH_CTX_free(s_context);
        s_context = NULL;
    }
    return all_free ? ESP_OK : ESP_ERR_TIMEOUT;
}

static void complete_lifecycle(uint32_t sequence, bool running,
                               esp_err_t result)
{
    taskENTER_CRITICAL(&s_lock);
    s_running = running;
    s_transitioning = false;
    s_completed_sequence = sequence;
    s_command_result = result;
    s_last_error = result;
    if (!running) {
        s_cleanup_pending = result != ESP_OK;
    }
    if (result == ESP_OK) {
        if (running) {
            ++s_counters.starts;
        } else {
            ++s_counters.stops;
        }
    } else if (s_desired_running) {
        ++s_counters.start_failures;
    }
    taskEXIT_CRITICAL(&s_lock);
}

static void handle_lifecycle(void)
{
    bool desired;
    bool running;
    bool transitioning;
    bool cleanup_pending;
    uint32_t sequence;
    taskENTER_CRITICAL(&s_lock);
    desired = s_desired_running;
    running = s_running;
    transitioning = s_transitioning;
    cleanup_pending = s_cleanup_pending;
    sequence = s_requested_sequence;
    taskEXIT_CRITICAL(&s_lock);

    if (!transitioning) {
        return;
    }
    if (desired == running && !(cleanup_pending && !desired)) {
        complete_lifecycle(sequence, running, ESP_OK);
        return;
    }

    if (desired) {
        esp_err_t error = start_runtime();
        complete_lifecycle(sequence, error == ESP_OK, error);
        if (error == ESP_OK) {
            ESP_LOGI(TAG, "SSH listening on TCP port %u", SSH_TRANSPORT_PORT);
        } else {
            ESP_LOGE(TAG, "SSH startup failed: %s", esp_err_to_name(error));
        }
    } else {
        esp_err_t error = stop_runtime();
        complete_lifecycle(sequence, false, error);
        if (error == ESP_OK) {
            ESP_LOGI(TAG, "SSH stopped");
        } else {
            ESP_LOGE(TAG, "SSH listener stopped but session cleanup failed: %s",
                     esp_err_to_name(error));
        }
    }
}

static ssh_slot_t *find_free_slot(size_t *slot_index)
{
    unsigned int pending_handshakes = 0U;
    for (size_t index = 0U; index < SSH_TRANSPORT_MAX_SESSIONS; ++index) {
        if (s_slots[index].state == SSH_TRANSPORT_SESSION_HANDSHAKE) {
            ++pending_handshakes;
        }
    }
    if (pending_handshakes >= SSH_TRANSPORT_MAX_PENDING_HANDSHAKES) {
        return NULL;
    }

    for (size_t index = 0U; index < SSH_TRANSPORT_MAX_SESSIONS; ++index) {
        if (s_slots[index].state == SSH_TRANSPORT_SESSION_FREE) {
            *slot_index = index;
            return &s_slots[index];
        }
    }
    return NULL;
}

static void format_peer(const struct sockaddr_storage *address,
                        char *output, size_t output_size)
{
    output[0] = '\0';
    if (address->ss_family == AF_INET) {
        const struct sockaddr_in *ipv4 = (const struct sockaddr_in *)address;
        char host[INET_ADDRSTRLEN] = {0};
        if (inet_ntop(AF_INET, &ipv4->sin_addr, host, sizeof(host)) != NULL) {
            (void)snprintf(output, output_size, "%s:%u", host,
                           (unsigned int)ntohs(ipv4->sin_port));
        }
    } else if (address->ss_family == AF_INET6) {
        const struct sockaddr_in6 *ipv6 = (const struct sockaddr_in6 *)address;
        char host[INET6_ADDRSTRLEN] = {0};
        if (inet_ntop(AF_INET6, &ipv6->sin6_addr, host, sizeof(host)) != NULL) {
            (void)snprintf(output, output_size, "[%s]:%u", host,
                           (unsigned int)ntohs(ipv6->sin6_port));
        }
    }
    if (output[0] == '\0') {
        strncpy(output, "unknown", output_size - 1U);
        output[output_size - 1U] = '\0';
    }
}

static void accept_connections(void)
{
    if (s_listen_fd < 0 || s_context == NULL) {
        return;
    }

    for (unsigned int accepted_count = 0U;
         accepted_count < SSH_TRANSPORT_ACCEPT_BUDGET;
         ++accepted_count) {
        struct sockaddr_storage peer_address;
        socklen_t peer_length = sizeof(peer_address);
        int socket_fd = accept(s_listen_fd, (struct sockaddr *)&peer_address,
                               &peer_length);
        if (socket_fd < 0) {
            if (errno != EAGAIN && errno != EWOULDBLOCK) {
                add_counter(&s_counters.io_failures, 1U);
            }
            return;
        }
        add_counter(&s_counters.tcp_connections, 1U);

        size_t slot_index = 0U;
        ssh_slot_t *slot = find_free_slot(&slot_index);
        if (slot == NULL) {
            add_counter(&s_counters.capacity_rejections, 1U);
            close(socket_fd);
            continue;
        }
        if (set_nonblocking(socket_fd) != ESP_OK) {
            add_counter(&s_counters.io_failures, 1U);
            close(socket_fd);
            continue;
        }
        int enabled = 1;
        (void)setsockopt(socket_fd, IPPROTO_TCP, TCP_NODELAY,
                         &enabled, sizeof(enabled));

        uint32_t generation = next_generation(slot->generation);
        memset(slot, 0, sizeof(*slot));
        slot->state = SSH_TRANSPORT_SESSION_HANDSHAKE;
        slot->generation = generation;
        slot->session_id = make_session_id(slot_index, generation);
        slot->socket_fd = socket_fd;
        slot->handshake_deadline_us = esp_timer_get_time() +
            (int64_t)SSH_TRANSPORT_HANDSHAKE_TIMEOUT_SECONDS * 1000000LL;
        format_peer(&peer_address, slot->peer, sizeof(slot->peer));

        slot->ssh = wolfSSH_new(s_context);
        if (slot->ssh == NULL || wolfSSH_set_fd(slot->ssh, socket_fd) != WS_SUCCESS) {
            add_counter(&s_counters.handshake_failures, 1U);
            (void)cleanup_slot(slot);
            publish_slot(slot, slot_index);
            continue;
        }
        wolfSSH_SetIOReadCtx(slot->ssh, slot);
        wolfSSH_SetUserAuthCtx(slot->ssh, slot);
        wolfSSH_SetUserAuthResultCtx(slot->ssh, slot);
        wolfSSH_SetChannelReqCtx(slot->ssh, slot);
        publish_slot(slot, slot_index);
    }
}

static bool slot_principal_is_current(const ssh_slot_t *slot)
{
    bool current = false;
    return slot->principal_valid &&
           user_database_principal_is_current(&slot->principal, &current) == ESP_OK &&
           current;
}

static void disconnect_failed_admission(ssh_slot_t *slot)
{
    if (slot->broker_client_id == SESSION_BROKER_NO_CLIENT) {
        return;
    }
    esp_err_t error = session_broker_disconnect(slot->broker_client_id);
    if (error == ESP_OK || error == ESP_ERR_NOT_FOUND) {
        slot->broker_client_id = SESSION_BROKER_NO_CLIENT;
        slot->writer = false;
        add_counter(&s_counters.disconnections, 1U);
    } else {
        add_counter(&s_counters.broker_failures, 1U);
    }
}

static esp_err_t connect_broker(ssh_slot_t *slot, size_t slot_index)
{
    if (!serial_service_is_running()) {
        esp_err_t error = serial_service_start();
        if (error != ESP_OK && serial_service_is_running()) {
            error = ESP_OK;
        }
        if (error != ESP_OK) {
            return error;
        }
    }

    if (!slot_principal_is_current(slot)) {
        return ESP_ERR_INVALID_STATE;
    }

    char name[SESSION_BROKER_CLIENT_NAME_MAX + 1U];
    int written = snprintf(name, sizeof(name), "ssh-%u-%" PRIu32,
                           (unsigned int)slot_index, slot->generation);
    if (written <= 0 || (size_t)written >= sizeof(name)) {
        return ESP_ERR_INVALID_SIZE;
    }

    esp_err_t error = session_broker_connect(
        SESSION_BROKER_CLIENT_SSH, name, &slot->broker_client_id);
    if (error != ESP_OK) {
        return error;
    }
    add_counter(&s_counters.broker_connections, 1U);
    if (!slot_principal_is_current(slot)) {
        disconnect_failed_admission(slot);
        return ESP_ERR_INVALID_STATE;
    }
    add_counter(&s_counters.writer_requests, 1U);
    error = session_broker_request_writer(slot->broker_client_id);
    if (error == ESP_OK) {
        slot->writer = true;
        add_counter(&s_counters.writer_grants, 1U);
    } else if (error == ESP_ERR_INVALID_STATE) {
        slot->writer = false;
        add_counter(&s_counters.writer_denials, 1U);
    } else {
        return error;
    }
    if (!slot_principal_is_current(slot)) {
        disconnect_failed_admission(slot);
        return ESP_ERR_INVALID_STATE;
    }
    return ESP_OK;
}

static void process_handshake(ssh_slot_t *slot, size_t slot_index)
{
    if (esp_timer_get_time() >= slot->handshake_deadline_us) {
        add_counter(&s_counters.handshake_timeouts, 1U);
        request_slot_close(slot, false);
        return;
    }

    slot->io_read_budget = SSH_TRANSPORT_WOLFSSH_READ_BUDGET;
    int result = wolfSSH_accept(slot->ssh);
    if (result == WS_SUCCESS) {
        bool principal_current = false;
        esp_err_t principal_error = slot->principal_valid
                                        ? user_database_principal_is_current(
                                              &slot->principal, &principal_current)
                                        : ESP_ERR_INVALID_STATE;
        if (!slot->authenticated || !slot->principal_valid) {
            request_slot_close(slot, false);
            return;
        }
        if (principal_error != ESP_OK || !principal_current) {
            request_slot_close(slot, true);
            return;
        }
        if (!slot->shell_requested ||
            wolfSSH_GetSessionType(slot->ssh) != WOLFSSH_SESSION_SHELL) {
            add_counter(&s_counters.request_rejections, 1U);
            request_slot_close(slot, false);
            return;
        }

        esp_err_t error;
        if (slot->principal.role == USER_ROLE_USER) {
            error = connect_broker(slot, slot_index);
            if (error != ESP_OK) {
                add_counter(&s_counters.broker_failures, 1U);
                request_slot_close(slot, false);
                return;
            }
            slot->route = SSH_TRANSPORT_ROUTE_BROKER;
        } else if (slot->principal.role == USER_ROLE_ADMIN) {
            admin_ssh_console_token_t token = admin_console_token(slot, slot_index);
            error = admin_ssh_console_open(&token, &slot->principal);
            if (error != ESP_OK) {
                add_counter(&s_counters.admin_console_admission_failures, 1U);
                request_slot_close(slot, false);
                return;
            }
            slot->route = SSH_TRANSPORT_ROUTE_ADMIN_CONSOLE;
            add_counter(&s_counters.admin_console_admissions, 1U);
        } else {
            request_slot_close(slot, true);
            return;
        }
        if (!slot_principal_is_current(slot)) {
            if (slot->route == SSH_TRANSPORT_ROUTE_BROKER) {
                disconnect_failed_admission(slot);
            }
            request_slot_close(slot, true);
            return;
        }
        slot->state = SSH_TRANSPORT_SESSION_ACTIVE;
        slot->last_reconcile_us = esp_timer_get_time();
        add_counter(&s_counters.handshake_successes, 1U);
        return;
    }
    if (!wolfssh_would_block(slot->ssh, result)) {
        add_counter(&s_counters.handshake_failures, 1U);
        request_slot_close(slot, false);
    }
}

static bool reconcile_writer(ssh_slot_t *slot)
{
    int64_t now = esp_timer_get_time();
    if (now - slot->last_reconcile_us < SSH_TRANSPORT_RECONCILE_INTERVAL_US) {
        return true;
    }
    slot->last_reconcile_us = now;

    bool principal_current = false;
    esp_err_t principal_error = slot->principal_valid
                                    ? user_database_principal_is_current(
                                          &slot->principal, &principal_current)
                                    : ESP_ERR_INVALID_STATE;
    if (principal_error != ESP_OK || !principal_current) {
        request_slot_close(slot, true);
        return true;
    }

    session_broker_client_snapshot_t snapshot;
    esp_err_t error = session_broker_get_client_snapshot(
        slot->broker_client_id, &snapshot);
    if (error == ESP_ERR_TIMEOUT) {
        return true;
    }
    if (error != ESP_OK) {
        return false;
    }
    if (slot->writer && !snapshot.is_writer) {
        slot->rx_offset = 0U;
        slot->rx_length = 0U;
        add_counter(&s_counters.writer_revocations, 1U);
    }
    slot->writer = snapshot.is_writer;
    return true;
}

static bool drain_broker_events(ssh_slot_t *slot)
{
    for (unsigned int count = 0U; count < 4U; ++count) {
        session_broker_event_t event;
        esp_err_t error = session_broker_pop_event(slot->broker_client_id, &event);
        if (error == ESP_ERR_TIMEOUT) {
            return true;
        }
        if (error != ESP_OK) {
            return false;
        }
        bool writer = event.writer_id == slot->broker_client_id;
        if (slot->writer && !writer) {
            slot->rx_offset = 0U;
            slot->rx_length = 0U;
            add_counter(&s_counters.writer_revocations, 1U);
        }
        slot->writer = writer;
    }
    return true;
}

static bool service_wolfssh_io(ssh_slot_t *slot)
{
    word32 channel_id = 0U;
    slot->io_read_budget = SSH_TRANSPORT_WOLFSSH_READ_BUDGET;
    int result = wolfSSH_worker(slot->ssh, &channel_id);
    (void)channel_id;
    if (result == WS_SUCCESS || wolfssh_would_block(slot->ssh, result)) {
        return true;
    }
    return false;
}

static bool flush_client_input(ssh_slot_t *slot)
{
    if (slot->rx_offset >= slot->rx_length) {
        slot->rx_offset = 0U;
        slot->rx_length = 0U;
        return true;
    }
    size_t remaining = slot->rx_length - slot->rx_offset;
    if (!slot->writer) {
        add_counter(&s_counters.rx_rejected_bytes, remaining);
        slot->rx_offset = 0U;
        slot->rx_length = 0U;
        return true;
    }

    size_t accepted = 0U;
    esp_err_t error = session_broker_write(
        slot->broker_client_id, slot->rx_buffer + slot->rx_offset,
        remaining, &accepted);
    if (error == ESP_ERR_TIMEOUT) {
        return true;
    }
    if (error == ESP_ERR_INVALID_STATE) {
        slot->writer = false;
        add_counter(&s_counters.rx_rejected_bytes, remaining);
        slot->rx_offset = 0U;
        slot->rx_length = 0U;
        return true;
    }
    if (error != ESP_OK) {
        return false;
    }
    slot->rx_offset += accepted;
    add_counter(&s_counters.rx_accepted_bytes, accepted);
    if (slot->rx_offset >= slot->rx_length) {
        slot->rx_offset = 0U;
        slot->rx_length = 0U;
    }
    return true;
}

static bool receive_client_input(ssh_slot_t *slot)
{
    if (slot->rx_length != 0U) {
        return true;
    }
    slot->io_read_budget = SSH_TRANSPORT_WOLFSSH_READ_BUDGET;
    int result = wolfSSH_stream_read(slot->ssh, slot->rx_buffer,
                                     sizeof(slot->rx_buffer));
    if (result > 0) {
        slot->rx_offset = 0U;
        slot->rx_length = (size_t)result;
        add_counter(&s_counters.rx_bytes, (uint64_t)result);
        return flush_client_input(slot);
    }
    if (result == 0 || wolfssh_would_block(slot->ssh, result)) {
        return true;
    }
    return false;
}

static bool flush_client_output(ssh_slot_t *slot)
{
    if (slot->tx_offset >= slot->tx_length) {
        slot->tx_offset = 0U;
        slot->tx_length = 0U;
        return true;
    }

    slot->io_read_budget = SSH_TRANSPORT_WOLFSSH_READ_BUDGET;
    int result = wolfSSH_stream_send(
        slot->ssh, slot->tx_buffer + slot->tx_offset,
        (word32)(slot->tx_length - slot->tx_offset));
    if (result > 0) {
        slot->tx_offset += (size_t)result;
        add_counter(&s_counters.tx_bytes, (uint64_t)result);
        if (slot->tx_offset >= slot->tx_length) {
            slot->tx_offset = 0U;
            slot->tx_length = 0U;
        }
        return true;
    }
    if (result == 0 || wolfssh_would_block(slot->ssh, result)) {
        return true;
    }
    return false;
}

static bool read_broker_output(ssh_slot_t *slot)
{
    if (slot->tx_length != 0U) {
        return true;
    }
    size_t received = 0U;
    esp_err_t error = session_broker_read(
        slot->broker_client_id, slot->tx_buffer,
        sizeof(slot->tx_buffer), &received);
    if (error == ESP_ERR_TIMEOUT) {
        return true;
    }
    if (error != ESP_OK) {
        return false;
    }
    slot->tx_offset = 0U;
    slot->tx_length = received;
    return received == 0U ? true : flush_client_output(slot);
}

static bool reconcile_admin_principal(ssh_slot_t *slot)
{
    int64_t now = esp_timer_get_time();
    if (now - slot->last_reconcile_us < SSH_TRANSPORT_RECONCILE_INTERVAL_US) {
        return true;
    }
    slot->last_reconcile_us = now;
    if (!slot_principal_is_current(slot) || slot->principal.role != USER_ROLE_ADMIN) {
        request_slot_close(slot, true);
        return false;
    }
    return true;
}

static bool flush_admin_input(ssh_slot_t *slot, size_t slot_index)
{
    if (slot->rx_offset >= slot->rx_length) {
        slot->rx_offset = 0U;
        slot->rx_length = 0U;
        return true;
    }
    admin_ssh_console_token_t token = admin_console_token(slot, slot_index);
    size_t consumed = 0U;
    bool accepted = admin_ssh_console_feed_input(
        &token, slot->rx_buffer + slot->rx_offset,
        slot->rx_length - slot->rx_offset, &consumed);
    if (consumed > 0U) {
        slot->rx_offset += consumed;
        add_counter(&s_counters.rx_accepted_bytes, consumed);
    }
    if (slot->rx_offset >= slot->rx_length) {
        slot->rx_offset = 0U;
        slot->rx_length = 0U;
    }
    if (!accepted && consumed == 0U) {
        add_counter(&s_counters.admin_console_input_rejections, 1U);
    }
    return true;
}

static bool receive_admin_input(ssh_slot_t *slot, size_t slot_index)
{
    if (slot->rx_length != 0U) {
        return flush_admin_input(slot, slot_index);
    }
    admin_ssh_console_token_t token = admin_console_token(slot, slot_index);
    if (!admin_ssh_console_accepts_input(&token)) {
        return true;
    }
    slot->io_read_budget = SSH_TRANSPORT_WOLFSSH_READ_BUDGET;
    int result = wolfSSH_stream_read(slot->ssh, slot->rx_buffer,
                                     sizeof(slot->rx_buffer));
    if (result > 0) {
        slot->rx_offset = 0U;
        slot->rx_length = (size_t)result;
        add_counter(&s_counters.rx_bytes, (uint64_t)result);
        return flush_admin_input(slot, slot_index);
    }
    return result == 0 || wolfssh_would_block(slot->ssh, result);
}

static bool read_admin_output(ssh_slot_t *slot, size_t slot_index)
{
    if (slot->tx_length != 0U) {
        return true;
    }
    admin_ssh_console_token_t token = admin_console_token(slot, slot_index);
    size_t received = 0U;
    esp_err_t error = admin_ssh_console_read_output(&token, slot->tx_buffer,
                                                    sizeof(slot->tx_buffer), &received);
    if (error != ESP_OK && error != ESP_ERR_NOT_FOUND) {
        return false;
    }
    slot->tx_offset = 0U;
    slot->tx_length = received;
    return error == ESP_OK;
}

static void process_active(ssh_slot_t *slot, size_t slot_index)
{
    bool healthy = service_wolfssh_io(slot);
    if (healthy && slot->route == SSH_TRANSPORT_ROUTE_BROKER) {
        healthy = drain_broker_events(slot) && reconcile_writer(slot) &&
                  flush_client_output(slot) && read_broker_output(slot) &&
                  flush_client_input(slot) && receive_client_input(slot);
    } else if (healthy && slot->route == SSH_TRANSPORT_ROUTE_ADMIN_CONSOLE) {
        healthy = reconcile_admin_principal(slot) && flush_client_output(slot) &&
                  read_admin_output(slot, slot_index) && flush_client_output(slot) &&
                  receive_admin_input(slot, slot_index);
    } else if (healthy) {
        healthy = false;
    }
    if (!healthy) {
        add_counter(&s_counters.io_failures, 1U);
        request_slot_close(slot, false);
    }
}

static void process_slots(void)
{
    for (size_t index = 0U; index < SSH_TRANSPORT_MAX_SESSIONS; ++index) {
        ssh_slot_t *slot = &s_slots[index];
        if (consume_external_close(slot, index)) {
            request_slot_close(slot, true);
        }
        if (slot->state == SSH_TRANSPORT_SESSION_FREE) {
            publish_slot(slot, index);
            continue;
        }
        if (slot->close_requested || slot->state == SSH_TRANSPORT_SESSION_CLOSING) {
            (void)cleanup_slot(slot);
            publish_slot(slot, index);
            continue;
        }
        if (slot->state == SSH_TRANSPORT_SESSION_HANDSHAKE) {
            process_handshake(slot, index);
        } else if (slot->state == SSH_TRANSPORT_SESSION_ACTIVE) {
            process_active(slot, index);
        }
        publish_slot(slot, index);
    }

    bool all_free = true;
    for (size_t index = 0U; index < SSH_TRANSPORT_MAX_SESSIONS; ++index) {
        if (s_slots[index].state != SSH_TRANSPORT_SESSION_FREE) {
            all_free = false;
            break;
        }
    }
    if (all_free) {
        taskENTER_CRITICAL(&s_lock);
        if (!s_running) {
            s_cleanup_pending = false;
        }
        taskEXIT_CRITICAL(&s_lock);
    }
}

static void transport_task(void *context)
{
    (void)context;
    for (;;) {
        handle_lifecycle();

        bool running;
        taskENTER_CRITICAL(&s_lock);
        running = s_running;
        taskEXIT_CRITICAL(&s_lock);
        if (running) {
            accept_connections();
        }
        process_slots();
        (void)ulTaskNotifyTake(pdTRUE, pdMS_TO_TICKS(SSH_TRANSPORT_LOOP_DELAY_MS));
    }
}

esp_err_t ssh_transport_init(void)
{
    taskENTER_CRITICAL(&s_lock);
    if (s_initialized) {
        taskEXIT_CRITICAL(&s_lock);
        return ESP_OK;
    }
    if (s_initializing || s_init_faulted) {
        taskEXIT_CRITICAL(&s_lock);
        return ESP_ERR_INVALID_STATE;
    }
    s_initializing = true;
    taskEXIT_CRITICAL(&s_lock);

    bool wolfssh_initialized = false;
    esp_err_t error = secure_random_init();
    if (error != ESP_OK) {
        error = ESP_ERR_INVALID_STATE;
        goto fail;
    }
    if (wolfSSL_SetAllocators(ssh_malloc, ssh_free, ssh_realloc) != 0) {
        error = ESP_FAIL;
        goto fail;
    }
    if (wolfSSH_Init() != WS_SUCCESS) {
        error = ESP_FAIL;
        goto fail;
    }
    wolfssh_initialized = true;
    if (wc_SetSeed_Cb(ssh_seed) != 0) {
        error = ESP_FAIL;
        goto fail;
    }

    SemaphoreHandle_t command_mutex = xSemaphoreCreateMutex();
    if (command_mutex == NULL) {
        error = ESP_ERR_NO_MEM;
        goto fail;
    }
    error = admin_ssh_console_init();
    if (error != ESP_OK) {
        vSemaphoreDelete(command_mutex);
        goto fail;
    }
    for (size_t index = 0U; index < SSH_TRANSPORT_MAX_SESSIONS; ++index) {
        s_slots[index].state = SSH_TRANSPORT_SESSION_FREE;
        s_slots[index].socket_fd = -1;
        s_session_snapshots[index].socket_fd = -1;
        s_session_snapshots[index].state = SSH_TRANSPORT_SESSION_FREE;
    }
    BaseType_t created = xTaskCreatePinnedToCore(
        transport_task, "ssh_transport", SSH_TRANSPORT_TASK_STACK_SIZE, NULL,
        SSH_TRANSPORT_TASK_PRIORITY, &s_task, SSH_TRANSPORT_TASK_CORE);
    if (created != pdPASS) {
        s_task = NULL;
        vSemaphoreDelete(command_mutex);
        error = ESP_ERR_NO_MEM;
        goto fail;
    }

    taskENTER_CRITICAL(&s_lock);
    s_command_mutex = command_mutex;
    s_initialized = true;
    s_initializing = false;
    s_last_error = ESP_OK;
    taskEXIT_CRITICAL(&s_lock);
    return ESP_OK;

fail:
    bool cleanup_failed = wolfssh_initialized && wolfSSH_Cleanup() != WS_SUCCESS;
    taskENTER_CRITICAL(&s_lock);
    s_init_faulted = cleanup_failed;
    s_initializing = false;
    taskEXIT_CRITICAL(&s_lock);
    return error;
}

static esp_err_t request_running_locked(bool desired)
{
    uint32_t sequence;
    taskENTER_CRITICAL(&s_lock);
    if (!s_initialized) {
        taskEXIT_CRITICAL(&s_lock);
        return ESP_ERR_INVALID_STATE;
    }
    if (s_transitioning) {
        taskEXIT_CRITICAL(&s_lock);
        return ESP_ERR_TIMEOUT;
    }
    if (desired && s_cleanup_pending) {
        taskEXIT_CRITICAL(&s_lock);
        return ESP_ERR_INVALID_STATE;
    }
    if (s_running == desired && !(s_cleanup_pending && !desired)) {
        taskEXIT_CRITICAL(&s_lock);
        return ESP_OK;
    }
    s_desired_running = desired;
    s_transitioning = true;
    s_requested_sequence = next_generation(s_requested_sequence);
    sequence = s_requested_sequence;
    taskEXIT_CRITICAL(&s_lock);
    notify_task();

    TickType_t deadline = xTaskGetTickCount() +
                          pdMS_TO_TICKS(SSH_TRANSPORT_COMMAND_TIMEOUT_MS);
    for (;;) {
        uint32_t completed;
        esp_err_t result;
        taskENTER_CRITICAL(&s_lock);
        completed = s_completed_sequence;
        result = s_command_result;
        taskEXIT_CRITICAL(&s_lock);
        if (completed == sequence) {
            return result;
        }
        if ((int32_t)(xTaskGetTickCount() - deadline) >= 0) {
            return ESP_ERR_TIMEOUT;
        }
        vTaskDelay(pdMS_TO_TICKS(10U));
    }
}

static esp_err_t request_running(bool desired)
{
    if (s_command_mutex == NULL) {
        return ESP_ERR_INVALID_STATE;
    }
    xSemaphoreTake(s_command_mutex, portMAX_DELAY);
    esp_err_t result = request_running_locked(desired);
    xSemaphoreGive(s_command_mutex);
    return result;
}

esp_err_t ssh_transport_start(void)
{
    return request_running(true);
}

esp_err_t ssh_transport_stop(void)
{
    return request_running(false);
}

esp_err_t ssh_transport_replace_host_key(bool reset)
{
    if (s_command_mutex == NULL) {
        return ESP_ERR_INVALID_STATE;
    }
    xSemaphoreTake(s_command_mutex, portMAX_DELAY);

    bool was_running;
    bool cleanup_pending;
    taskENTER_CRITICAL(&s_lock);
    if (!s_initialized || s_transitioning) {
        taskEXIT_CRITICAL(&s_lock);
        xSemaphoreGive(s_command_mutex);
        return ESP_ERR_INVALID_STATE;
    }
    was_running = s_running;
    cleanup_pending = s_cleanup_pending;
    taskEXIT_CRITICAL(&s_lock);

    esp_err_t error = ESP_OK;
    if (was_running || cleanup_pending) {
        error = request_running_locked(false);
    }
    if (error == ESP_OK) {
        error = reset ? ssh_security_reset() : ssh_security_rotate();
    }
    if (error != ESP_OK) {
        if (was_running) {
            (void)request_running_locked(true);
        }
        xSemaphoreGive(s_command_mutex);
        return error;
    }
    if (was_running || reset) {
        error = request_running_locked(true);
    }
    xSemaphoreGive(s_command_mutex);
    return error;
}

esp_err_t ssh_transport_get_snapshot(ssh_transport_snapshot_t *snapshot)
{
    if (snapshot == NULL) {
        return ESP_ERR_INVALID_ARG;
    }

    taskENTER_CRITICAL(&s_lock);
    if (!s_initialized) {
        taskEXIT_CRITICAL(&s_lock);
        return ESP_ERR_INVALID_STATE;
    }
    TaskHandle_t task = s_task;
    memset(snapshot, 0, sizeof(*snapshot));
    snapshot->initialized = s_initialized;
    snapshot->running = s_running;
    snapshot->transitioning = s_transitioning;
    snapshot->port = SSH_TRANSPORT_PORT;
    snapshot->last_error = s_last_error;
    snapshot->task_core_id = SSH_TRANSPORT_TASK_CORE;
    snapshot->task_stack_size = SSH_TRANSPORT_TASK_STACK_SIZE;
    snapshot->counters = s_counters;
    for (size_t index = 0U; index < SSH_TRANSPORT_MAX_SESSIONS; ++index) {
        snapshot->sessions[index] = s_session_snapshots[index];
        if (snapshot->sessions[index].active) {
            ++snapshot->active_sessions;
        }
    }
    taskEXIT_CRITICAL(&s_lock);
    if (task != NULL) {
        snapshot->task_stack_free_minimum =
            (uint32_t)uxTaskGetStackHighWaterMark(task);
    }
    return ESP_OK;
}

esp_err_t ssh_transport_clear_counters(void)
{
    taskENTER_CRITICAL(&s_lock);
    if (!s_initialized) {
        taskEXIT_CRITICAL(&s_lock);
        return ESP_ERR_INVALID_STATE;
    }
    memset(&s_counters, 0, sizeof(s_counters));
    taskEXIT_CRITICAL(&s_lock);
    return ESP_OK;
}

esp_err_t ssh_transport_disconnect(uint32_t session_id)
{
    if (session_id == 0U) {
        return ESP_ERR_INVALID_ARG;
    }

    bool found = false;
    taskENTER_CRITICAL(&s_lock);
    if (!s_initialized) {
        taskEXIT_CRITICAL(&s_lock);
        return ESP_ERR_INVALID_STATE;
    }
    for (size_t index = 0U; index < SSH_TRANSPORT_MAX_SESSIONS; ++index) {
        if (s_session_snapshots[index].active &&
            s_session_snapshots[index].session_id == session_id) {
            s_external_close_id[index] = session_id;
            found = true;
            break;
        }
    }
    taskEXIT_CRITICAL(&s_lock);
    if (found) {
        notify_task();
        return ESP_OK;
    }
    return ESP_ERR_NOT_FOUND;
}

esp_err_t ssh_transport_revoke_user(const uint8_t *username,
                                    size_t username_length)
{
    if (!user_database_username_valid(username, username_length)) {
        return ESP_ERR_INVALID_ARG;
    }

    bool found = false;
    taskENTER_CRITICAL(&s_lock);
    if (!s_initialized) {
        taskEXIT_CRITICAL(&s_lock);
        return ESP_ERR_INVALID_STATE;
    }
    for (size_t index = 0U; index < SSH_TRANSPORT_MAX_SESSIONS; ++index) {
        const ssh_transport_session_snapshot_t *session =
            &s_session_snapshots[index];
        if (session->active && session->principal_valid &&
            strlen(session->username) == username_length &&
            memcmp(session->username, username, username_length) == 0) {
            s_external_close_id[index] = session->session_id;
            found = true;
        }
    }
    taskEXIT_CRITICAL(&s_lock);
    if (found) {
        notify_task();
    }
    return ESP_OK;
}

esp_err_t ssh_transport_revoke_sessions(void)
{
    taskENTER_CRITICAL(&s_lock);
    if (!s_initialized) {
        taskEXIT_CRITICAL(&s_lock);
        return ESP_ERR_INVALID_STATE;
    }
    for (size_t index = 0U; index < SSH_TRANSPORT_MAX_SESSIONS; ++index) {
        if (s_session_snapshots[index].active) {
            s_external_close_id[index] =
                s_session_snapshots[index].session_id;
        }
    }
    taskEXIT_CRITICAL(&s_lock);
    notify_task();
    return ESP_OK;
}