#ifndef ZK_AUTH_H
#define ZK_AUTH_H

#include "cryptoauthlib.h"
#include "host/atca_host.h"
#include "libssh2_config.h"
#include "mbedtls/sha256.h"
#include "netdb.h"
#include <cJSON.h>
#include <esp_mac.h>
#include <esp_system.h>
#include <libssh2.h>
#include <mbedtls/aes.h>
#include <mbedtls/base64.h>
#include <mbedtls/ctr_drbg.h>
#include <mbedtls/ecdh.h>
#include <mbedtls/ecp.h>
#include <mbedtls/entropy.h>
#include <mbedtls/md.h>
#include <mbedtls/pkcs5.h>
#include <mbedtls/sha256.h>
#include <stdio.h>
#include <string.h>
// Added network headers for ESP-IDF (lwIP) sockets
#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>

#define BUFSIZE 3200

class ZKAuth {
private:
  mbedtls_ecp_group grp;
  mbedtls_mpi device_private_d;
  mbedtls_ecp_point device_public_Q;
  mbedtls_entropy_context entropy;
  mbedtls_ctr_drbg_context ctr_drbg;

  uint8_t device_public_key[65];    // Uncompressed: 0x04 + X(32) + Y(32)
  uint8_t stored_password_hash[32]; // PBKDF2 hash of the correct password

  bool initialized;

  // Convert binary to hex string
  void bin_to_hex(const uint8_t *bin, size_t bin_len, char *hex) {
    for (size_t i = 0; i < bin_len; i++) {
      sprintf(hex + (i * 2), "%02x", bin[i]);
    }
    hex[bin_len * 2] = '\0';
  }

  // Convert hex string to binary
  bool hex_to_bin(const char *hex, uint8_t *bin, size_t bin_len) {
    if (strlen(hex) != bin_len * 2)
      return false;

    for (size_t i = 0; i < bin_len; i++) {
      if (sscanf(hex + (i * 2), "%2hhx", &bin[i]) != 1) {
        return false;
      }
    }
    return true;
  }

public:
  ZKAuth() : initialized(false) {
    mbedtls_ecp_group_init(&grp);
    mbedtls_mpi_init(&device_private_d);
    mbedtls_ecp_point_init(&device_public_Q);
    mbedtls_entropy_init(&entropy);
    mbedtls_ctr_drbg_init(&ctr_drbg);
  }

  ~ZKAuth() {
    mbedtls_ecp_group_free(&grp);
    mbedtls_mpi_free(&device_private_d);
    mbedtls_ecp_point_free(&device_public_Q);
    mbedtls_entropy_free(&entropy);
    mbedtls_ctr_drbg_free(&ctr_drbg);
  }

  // Initialize the ZK authentication system
  bool init() {
    if (initialized) {
      return true;
    } else {
      initialized = true;
      printf("\n=== ZK Authentication Initialized ===\n");
    }
    return true;
  }

  // Marked as static so it can be passed as a C callback to libssh2
  // Helper to safely format SSH integers (mpint) exactly how OpenSSH demands
  static uint32_t write_mpint(uint8_t *buf, const uint8_t *val, uint32_t size) {
    uint32_t start = 0;

    // Strip unnecessary leading zeros
    while (start < (size - 1) && val[start] == 0x00) {
      start++;
    }

    // If the most significant bit is 1, we must prepend a 0x00 padding byte
    bool pad = (val[start] & 0x80) != 0;
    uint32_t len = (size - start) + (pad ? 1 : 0);

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

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

  // The main signing callback
  static int sign_callback(LIBSSH2_SESSION *session, unsigned char **sig,
                           size_t *sig_len, const unsigned char *data,
                           size_t data_len, void **abstract) {
    uint8_t digest[32];
    uint8_t raw_sig[64];

    // 1. Hash the fully assembled challenge buffer provided by libssh2
    mbedtls_sha256(data, data_len, digest, 0);

    // 2. Request signature from the ATECC608B
    if (atcab_sign(0x0, digest, raw_sig) != ATCA_SUCCESS) {
      printf("CRITICAL: ATECC608B Signing Failed!\n");
      return -1;
    }

    // 3. Allocate memory JUST for the two mathematical integers (max ~74 bytes)
    unsigned char *buf = (unsigned char *)malloc(80);
    if (!buf)
      return -1;

    // 4. Format strictly as [mpint R] [mpint S]. NO outer strings!
    uint32_t offset = 0;
    offset += write_mpint(&buf[offset], &raw_sig[0], 32);  // Format R
    offset += write_mpint(&buf[offset], &raw_sig[32], 32); // Format S

    // Hand ownership to libssh2
    *sig = buf;
    *sig_len = offset;

    return 0; // Success!
  }

  // Helper function to write a 32-bit integer in big-endian format
  void write_uint32_be(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;
  }

  // Function to generate the authorized_keys string
  void generate_ssh_authorized_key(const uint8_t *atecc_pubkey,
                                   uint8_t *out_blob) {
    uint8_t ssh_blob[104];
    uint32_t offset = 0;

    // 1. Key Type
    const char *key_type = "ecdsa-sha2-nistp256";
    write_uint32_be(&ssh_blob[offset], 19);
    offset += 4;
    memcpy(&ssh_blob[offset], key_type, 19);
    offset += 19;

    // 2. Curve Name
    const char *curve_name = "nistp256";
    write_uint32_be(&ssh_blob[offset], 8);
    offset += 4;
    memcpy(&ssh_blob[offset], curve_name, 8);
    offset += 8;

    // 3. Public Key (Uncompressed format: 0x04 + 64 bytes X/Y)
    write_uint32_be(&ssh_blob[offset], 65);
    offset += 4;
    ssh_blob[offset++] = 0x04;
    memcpy(&ssh_blob[offset], atecc_pubkey, 64);
    offset += 64;

    // 4. Base64 Encode the blob
    size_t b64_len = 0;
    // Call once to get required length
    mbedtls_base64_encode(NULL, 0, &b64_len, ssh_blob, 104);

    unsigned char b64_out[b64_len];
    // Call again to actually encode
    mbedtls_base64_encode(b64_out, b64_len, &b64_len, ssh_blob, 104);

    // 5. Print out the final authorized_keys line
    printf("ecdsa-sha2-nistp256 %s esp32-atecc608b\n", b64_out);

    memcpy(out_blob, ssh_blob, 104);
  }

  static void my_trace_handler(LIBSSH2_SESSION *session, void *context,
                               const char *data, size_t length) {
    // Print the trace data to the console
    printf("LIBSSH2_TRACE: %.*s", (int)length, data);
  }

  static int waitsocket(int socket_fd, LIBSSH2_SESSION *session) {
    struct timeval timeout;
    int rc;
    fd_set fd;
    fd_set *writefd = NULL;
    fd_set *readfd = NULL;
    int dir;

    timeout.tv_sec = 10;
    timeout.tv_usec = 0;

    FD_ZERO(&fd);

    FD_SET(socket_fd, &fd);

    /* now make sure we wait in the correct direction */
    dir = libssh2_session_block_directions(session);

    if (dir & LIBSSH2_SESSION_BLOCK_INBOUND)
      readfd = &fd;

    if (dir & LIBSSH2_SESSION_BLOCK_OUTBOUND)
      writefd = &fd;

    rc = select(socket_fd + 1, readfd, writefd, NULL, &timeout);

    return rc;
  }

  // Get device identity (for /api/identity endpoint)
  char *get_identity_json() {
    char pubkey_hex[131];     // 65 bytes * 2 + null
    uint8_t atecc_pubkey[64]; // 65 bytes * 2 + null
    uint8_t standard_pubkey[65];
    standard_pubkey[0] = 0x04;
    // Get public key from ATECC608B and convert to hex
    ATCA_STATUS status = atcab_get_pubkey(0, atecc_pubkey);
    if (status != ATCA_SUCCESS) {
      printf("Failed to read public key from ATECC608B: 0x%02X\n", status);
    }

    uint8_t pubkey_blob[104];

    generate_ssh_authorized_key(atecc_pubkey, pubkey_blob);

    // Print the authorized_keys line for debugging
    printf("Generated authorized_keys line:\n");
    for (int i = 0; i < 104; i++) {
      printf("%02x", pubkey_blob[i]);
    }
    printf("\n");

    memcpy(&standard_pubkey[1], atecc_pubkey, 64);
    bin_to_hex(standard_pubkey, 65, pubkey_hex);

    // Get MAC address to use as salt
    uint8_t mac[6];
    esp_read_mac(mac, ESP_MAC_WIFI_STA);
    char mac_hex[13]; // 6 bytes * 2 + null
    bin_to_hex(mac, 6, mac_hex);

    cJSON *root = cJSON_CreateObject();
    cJSON_AddStringToObject(root, "pubKey", pubkey_hex);
    cJSON_AddStringToObject(root, "macAddress", mac_hex);

    char *json_str = cJSON_PrintUnformatted(root);
    cJSON_Delete(root);

    // Print the JSON for debugging
    printf("Identity JSON: %s\n", json_str);

    const char *username = "jonathan";

    int sock;
    struct sockaddr_in sin;
    const char *hostname = "192.168.178.86";
    unsigned short port = 22;
    void *my_abstract = NULL;
    LIBSSH2_SESSION *session;
    LIBSSH2_CHANNEL *channel;

    ESP_LOGI(TAG, "libssh2_version is %s", LIBSSH2_VERSION);
    int rc = libssh2_init(0);
    if (rc) {
      ESP_LOGE(TAG, "libssh2 initialization failed (%d)", rc);
      while (1) {
        vTaskDelay(1);
      }
    }

    ESP_LOGD(TAG, "hostname=%s", hostname);
    ESP_LOGD(TAG, "port=%d", port);
    sin.sin_family = AF_INET;
    // sin.sin_port = htons(22);
    sin.sin_port = htons(port);
    sin.sin_addr.s_addr = inet_addr(hostname);
    ESP_LOGD(TAG, "sin.sin_addr.s_addr=%" PRIx32, sin.sin_addr.s_addr);
    if (sin.sin_addr.s_addr == 0xffffffff) {
      struct hostent *hp;
      hp = gethostbyname(hostname);
      if (hp == NULL) {
        ESP_LOGE(TAG, "gethostbyname fail");
        ESP_LOGE(TAG, "hostname=%s", hostname);
        ESP_LOGE(TAG, "port=%d", port);
        while (1) {
          vTaskDelay(1);
        }
      }
      struct ip4_addr *ip4_addr;
      ip4_addr = (struct ip4_addr *)hp->h_addr;
      sin.sin_addr.s_addr = ip4_addr->addr;
      ESP_LOGD(TAG, "sin.sin_addr.s_addr=%" PRIx32, sin.sin_addr.s_addr);
    }

    sock = socket(AF_INET, SOCK_STREAM, 0);
    if (sock == -1) {
      ESP_LOGE(TAG, "failed to create socket!");
      while (1) {
        vTaskDelay(1);
      }
    }

    if (connect(sock, (struct sockaddr *)(&sin), sizeof(struct sockaddr_in)) !=
        0) {
      ESP_LOGE(TAG, "failed to connect!");
      ESP_LOGE(TAG, "hostname=%s", hostname);
      ESP_LOGE(TAG, "port=%d", port);
      while (1) {
        vTaskDelay(1);
      }
    }

    printf("Connected. Starting SSH handshake...\n");

    session = libssh2_session_init();
    if (!session) {
      ESP_LOGE(TAG, "failed to session init");
      while (1) {
        vTaskDelay(1);
      }
    }

    libssh2_trace(session, ~0);
    libssh2_trace_sethandler(session, NULL, my_trace_handler);

    printf("Session initialized. Setting timeout...\n");
    // libssh2_session_set_timeout(session, 10000);

    printf("Performing handshake...\n");
    rc = libssh2_session_handshake(session, sock);
    if (rc) {
      ESP_LOGE(TAG, "Failure establishing SSH session: %d", rc);
      while (1) {
        vTaskDelay(1);
      }
    }

    printf("Handshake completed. Authenticating with public key...\n");

    // PUBLIC KEY AUTH: using custom callback
    rc = libssh2_userauth_publickey(session, username, pubkey_blob,
                                    sizeof(pubkey_blob), sign_callback,
                                    &my_abstract);

    if (rc == 0) {
      printf("SSH authentication successful!\n");
      // ... Continue SSH communication ...
    } else {
      printf("Authentication failed\n");
    }

    // Cleanup
    libssh2_session_disconnect(session, "Bye");
    libssh2_session_free(session);
    close(sock);
    libssh2_exit();

    return json_str;
  }
};

#endif // ZK_AUTH_H