#include "ssh_client.h"

#include "cryptoauthlib.h"
#include "sdkconfig.h"
#include <arpa/inet.h>
#include <esp_log.h>
#include <libssh2.h>
#include <mbedtls/base64.h>
#include <mbedtls/sha256.h>
#include <netdb.h>
#include <netinet/in.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <unistd.h>

static const char *TAG = "ssh_client";

/* -------------------------------------------------------------------------
 * SSH blob / key helpers
 * ---------------------------------------------------------------------- */

/* Write a big-endian uint32 into buf. */
static void write_be32(uint8_t *buf, uint32_t val)
{
    buf[0] = (val >> 24) & 0xFF;
    buf[1] = (val >> 16) & 0xFF;
    buf[2] = (val >> 8)  & 0xFF;
    buf[3] =  val        & 0xFF;
}

/**
 * Format one ECDSA integer as an SSH mpint:
 *   4-byte big-endian length  |  optional 0x00 pad  |  value (leading zeros stripped)
 *
 * Returns number of bytes written to buf.
 */
static uint32_t write_mpint(uint8_t *buf, const uint8_t *val, uint32_t size)
{
    uint32_t start = 0;
    while (start < (size - 1) && val[start] == 0x00) {
        start++;
    }

    bool pad = (val[start] & 0x80) != 0;
    uint32_t len = (size - start) + (pad ? 1 : 0);

    uint32_t off = 0;
    buf[off++] = (len >> 24) & 0xFF;
    buf[off++] = (len >> 16) & 0xFF;
    buf[off++] = (len >> 8)  & 0xFF;
    buf[off++] =  len        & 0xFF;

    if (pad) {
        buf[off++] = 0x00;
    }
    memcpy(&buf[off], &val[start], size - start);
    return off + (size - start);
}

/**
 * Build a 104-byte SSH public key blob for ecdsa-sha2-nistp256:
 *   [uint32 19] "ecdsa-sha2-nistp256"
 *   [uint32  8] "nistp256"
 *   [uint32 65] 0x04 || X(32) || Y(32)
 *
 * @param atecc_pubkey  64-byte raw public key (X||Y) from ATECC608B.
 * @param out_blob      Must point to a buffer of at least 104 bytes.
 */
static void build_pubkey_blob(const uint8_t *atecc_pubkey, uint8_t *out_blob)
{
    uint32_t off = 0;

    /* Key type */
    const char *key_type = "ecdsa-sha2-nistp256";
    write_be32(&out_blob[off], 19);  off += 4;
    memcpy(&out_blob[off], key_type, 19); off += 19;

    /* Curve name */
    const char *curve = "nistp256";
    write_be32(&out_blob[off], 8);   off += 4;
    memcpy(&out_blob[off], curve, 8); off += 8;

    /* Uncompressed EC point: 0x04 || X || Y */
    write_be32(&out_blob[off], 65);  off += 4;
    out_blob[off++] = 0x04;
    memcpy(&out_blob[off], atecc_pubkey, 64);
    /* off += 64 — total = 104 */
}

/* -------------------------------------------------------------------------
 * libssh2 signing callback (ATECC608B signs the hash)
 * ---------------------------------------------------------------------- */
static int sign_callback(LIBSSH2_SESSION *session,
                         unsigned char **sig, size_t *sig_len,
                         const unsigned char *data, size_t data_len,
                         void **abstract)
{
    (void)session;
    (void)abstract;

    uint8_t digest[32];
    uint8_t raw_sig[64];

    /* Hash the challenge data */
    mbedtls_sha256(data, data_len, digest, 0);

    /* Sign with the ATECC608B hardware key in slot 0 */
    if (atcab_sign(0, digest, raw_sig) != ATCA_SUCCESS) {
        ESP_LOGE(TAG, "ATECC608B signing failed!");
        return -1;
    }

    /* Encode as [mpint R][mpint S] (no outer string wrapper) */
    unsigned char *buf = (unsigned char *)malloc(80);
    if (!buf) return -1;

    uint32_t off = 0;
    off += write_mpint(&buf[off], &raw_sig[0],  32); /* R */
    off += write_mpint(&buf[off], &raw_sig[32], 32); /* S */

    *sig     = buf;
    *sig_len = off;
    return 0;
}

/* -------------------------------------------------------------------------
 * Helpers for non-blocking libssh2
 * ---------------------------------------------------------------------- */
static int waitsocket(int sock_fd, LIBSSH2_SESSION *session)
{
    struct timeval tv = { .tv_sec = 10, .tv_usec = 0 };
    fd_set fd;
    FD_ZERO(&fd);
    FD_SET(sock_fd, &fd);

    fd_set *rfd = NULL, *wfd = NULL;
    int dir = libssh2_session_block_directions(session);
    if (dir & LIBSSH2_SESSION_BLOCK_INBOUND)  rfd = &fd;
    if (dir & LIBSSH2_SESSION_BLOCK_OUTBOUND) wfd = &fd;

    return select(sock_fd + 1, rfd, wfd, NULL, &tv);
}

/* -------------------------------------------------------------------------
 * Public API
 * ---------------------------------------------------------------------- */
void ssh_print_public_key(void)
{
    uint8_t raw_key[64];
    ATCA_STATUS st = atcab_get_pubkey(0, raw_key);
    if (st != ATCA_SUCCESS) {
        ESP_LOGE(TAG, "atcab_get_pubkey failed: 0x%02X", st);
        return;
    }

    uint8_t blob[104];
    build_pubkey_blob(raw_key, blob);

    size_t b64_len = 0;
    mbedtls_base64_encode(NULL, 0, &b64_len, blob, sizeof(blob));

    unsigned char *b64 = (unsigned char *)malloc(b64_len + 1);
    if (!b64) {
        ESP_LOGE(TAG, "malloc failed for base64 buffer");
        return;
    }
    mbedtls_base64_encode(b64, b64_len, &b64_len, blob, sizeof(blob));
    b64[b64_len] = '\0';

    printf("\n=== SSH Public Key (add to authorized_keys) ===\n");
    printf("ecdsa-sha2-nistp256 %s keypitecc\n", b64);
    printf("================================================\n\n");

    free(b64);
}

bool ssh_execute_command(const char *cmd)
{
    if (!cmd) return false;

    /* --- Read public key blob ------------------------------------------ */
    uint8_t raw_key[64];
    ATCA_STATUS st = atcab_get_pubkey(0, raw_key);
    if (st != ATCA_SUCCESS) {
        ESP_LOGE(TAG, "atcab_get_pubkey failed: 0x%02X", st);
        return false;
    }
    uint8_t pubkey_blob[104];
    build_pubkey_blob(raw_key, pubkey_blob);

    /* --- TCP connect ------------------------------------------------------- */
    int rc;
    int sock = socket(AF_INET, SOCK_STREAM, 0);
    if (sock < 0) {
        ESP_LOGE(TAG, "socket() failed");
        return false;
    }

    struct sockaddr_in sin = {};
    sin.sin_family = AF_INET;
    sin.sin_port   = htons(CONFIG_SSH_PORT);
    sin.sin_addr.s_addr = inet_addr(CONFIG_SSH_HOSTNAME);

    if (sin.sin_addr.s_addr == 0xFFFFFFFF) {
        struct hostent *hp = gethostbyname(CONFIG_SSH_HOSTNAME);
        if (!hp) {
            ESP_LOGE(TAG, "gethostbyname(%s) failed", CONFIG_SSH_HOSTNAME);
            close(sock);
            return false;
        }
        sin.sin_addr.s_addr = ((struct ip4_addr *)hp->h_addr)->addr;
    }

    if (connect(sock, (struct sockaddr *)&sin, sizeof(sin)) != 0) {
        ESP_LOGE(TAG, "connect to %s:%d failed", CONFIG_SSH_HOSTNAME, CONFIG_SSH_PORT);
        close(sock);
        return false;
    }
    ESP_LOGI(TAG, "TCP connected to %s:%d", CONFIG_SSH_HOSTNAME, CONFIG_SSH_PORT);

    /* --- libssh2 session -------------------------------------------------- */
    rc = libssh2_init(0);
    if (rc) {
        ESP_LOGE(TAG, "libssh2_init failed: %d", rc);
        close(sock);
        return false;
    }

    LIBSSH2_SESSION *session = libssh2_session_init();
    if (!session) {
        ESP_LOGE(TAG, "libssh2_session_init failed");
        libssh2_exit();
        close(sock);
        return false;
    }

    rc = libssh2_session_handshake(session, sock);
    if (rc) {
        ESP_LOGE(TAG, "SSH handshake failed: %d", rc);
        goto cleanup_session;
    }
    ESP_LOGI(TAG, "SSH handshake OK");

    /* --- Authenticate with ATECC608B hardware key ------------------------- */
    void *abstract = NULL;
    rc = libssh2_userauth_publickey(session, CONFIG_SSH_USERNAME,
                                    pubkey_blob, sizeof(pubkey_blob),
                                    sign_callback, &abstract);
    if (rc != 0) {
        ESP_LOGE(TAG, "SSH public-key authentication failed: %d", rc);
        goto cleanup_session;
    }
    ESP_LOGI(TAG, "SSH authenticated as '%s'", CONFIG_SSH_USERNAME);

    /* --- Open channel and execute command --------------------------------- */
    LIBSSH2_CHANNEL *channel;
    do {
        channel = libssh2_channel_open_session(session);
        if (!channel && libssh2_session_last_errno(session) == LIBSSH2_ERROR_EAGAIN) {
            waitsocket(sock, session);
        }
    } while (!channel);

    if (!channel) {
        ESP_LOGE(TAG, "Failed to open SSH channel");
        goto cleanup_session;
    }

    ESP_LOGI(TAG, "Executing: %s", cmd);
    while ((rc = libssh2_channel_exec(channel, cmd)) == LIBSSH2_ERROR_EAGAIN) {
        waitsocket(sock, session);
    }
    if (rc != 0) {
        ESP_LOGE(TAG, "channel_exec failed: %d", rc);
        libssh2_channel_free(channel);
        goto cleanup_session;
    }

    /* --- Read output ------------------------------------------------------- */
    char buf[256];
    int bytes;
    ESP_LOGI(TAG, "--- command output ---");
    for (;;) {
        do {
            bytes = libssh2_channel_read(channel, buf, sizeof(buf) - 1);
        } while (bytes == LIBSSH2_ERROR_EAGAIN && (waitsocket(sock, session), 1));

        if (bytes <= 0) break;
        buf[bytes] = '\0';
        printf("%s", buf);
    }
    ESP_LOGI(TAG, "--- end output ---");

    /* --- Shutdown --------------------------------------------------------- */
    while ((rc = libssh2_channel_send_eof(channel)) == LIBSSH2_ERROR_EAGAIN) {
        waitsocket(sock, session);
    }
    while (libssh2_channel_wait_eof(channel) == LIBSSH2_ERROR_EAGAIN) {
        waitsocket(sock, session);
    }
    while (libssh2_channel_wait_closed(channel) == LIBSSH2_ERROR_EAGAIN) {
        waitsocket(sock, session);
    }

    int exit_status = libssh2_channel_get_exit_status(channel);
    ESP_LOGI(TAG, "Command exited with status: %d", exit_status);

    libssh2_channel_free(channel);

    libssh2_session_disconnect(session, "Bye");
    libssh2_session_free(session);
    libssh2_exit();
    close(sock);
    return (exit_status == 0);

cleanup_session:
    libssh2_session_disconnect(session, "Error");
    libssh2_session_free(session);
    libssh2_exit();
    close(sock);
    return false;
}