#ifndef ZK_AUTH_H
#define ZK_AUTH_H

#include "cryptoauthlib.h"
#include "host/atca_host.h"
#include "mbedtls/sha256.h"
#include <cJSON.h>
#include <esp_mac.h>
#include <esp_system.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>

// Zero-Knowledge Authentication Module
// Implements Client-Side KDF + ECIES Tunnel for ESP32-C6
// Format: SHA-256 compatible with ATECC608B
//
// Encryption: AES-256-CBC + HMAC-SHA256 (Encrypt-then-MAC)
// - Provides confidentiality (CBC) + authenticity (HMAC)
// - Blob format: IV (16 bytes) + Ciphertext (32 bytes) + HMAC (32 bytes) = 80
// bytes
// - HMAC verified BEFORE decryption (prevents padding oracle attacks)

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;
  }

  // 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 *atec_pubkey) {
    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], atec_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);
  }

  // Get device identity (for /api/identity endpoint)
  char *get_identity_json() {
    char pubkey_hex[131];    // 65 bytes * 2 + null
    uint8_t atec_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, atec_pubkey);
    if (status != ATCA_SUCCESS) {
      printf("Failed to read public key from ATECC608B: 0x%02X\n", status);
    }

    generate_ssh_authorized_key(atec_pubkey);

    memcpy(&standard_pubkey[1], atec_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);
    return json_str;
  }
};

#endif // ZK_AUTH_H