/*
 * atcmd_tls.c
 */
#define MBEDTLS_CONFIG_FILE "config-hsm.h"
#include "lwip/apps/sntp.h"

#include <stdio.h>
#include <string.h>

#include "lwip/netdb.h"
#include "lwip/sockets.h"


#include "mbedtls/ctr_drbg.h"
#include "mbedtls/entropy.h"
#include "mbedtls/error.h"
#include "mbedtls/pk.h"
#include "mbedtls/ssl.h"
#include "mbedtls/x509_crt.h"
#include "psa/crypto.h"


#include "atcmd_tls.h"
// #include "cert_file.h"       /* ca_cert_pem, client_cert_pem, client_key_pem // */
#include "uart_term.h"
#include <mbedtls/debug.h>

#include "date_time_service.h"
#include "mbedtls/platform_time.h"
#include "sntp_task.h"
#include <time.h>


#include "client.h"
#include "client_key.h"
#include "mosquitto.org.h"
#include "sntp_wrapper.h"
#define WIFI_PLATFORM_MBEDTLS_DEBUG_LOG_LEVEL 10
#define MQTTS_SOCKET_TIMEOUT_MS 15000
/* ── Global state (accessed by mqtt_sl.c) ───────────────────── */
mbedtls_ssl_context ssl;
int sock_fd = -1;


/* ── Private state ───────────────────────────────────────────── */
static mbedtls_ssl_config conf;
static mbedtls_x509_crt ca_cert;
static mbedtls_x509_crt client_cert;
static mbedtls_pk_context client_key;
// static mbedtls_entropy_context  entropy;
// static mbedtls_ctr_drbg_context ctr_drbg;

/* PSA RNG wrapper */
static int psa_rng_wrapper(void *ctx, unsigned char *output, size_t len) {
  (void)ctx;
  psa_status_t status = psa_generate_random(output, len);
  return (status == PSA_SUCCESS) ? 0 : -1;
}

/* ── BIO callbacks ───────────────────────────────────────────── */

static int lwip_bio_send(void *ctx, const unsigned char *buf, size_t len) {
  int fd = *((int *)ctx);
  int ret = lwip_send(fd, buf, len, 0);
  if (ret < 0) {
    return MBEDTLS_ERR_SSL_INTERNAL_ERROR;
  }
  return ret;
}

static int lwip_bio_recv(void *ctx, unsigned char *buf, size_t len) {
  int fd = *((int *)ctx);
  int ret = lwip_recv(fd, buf, len, 0);
  if (ret < 0) {
    return MBEDTLS_ERR_SSL_INTERNAL_ERROR;
  }
  if (ret == 0) {
    return MBEDTLS_ERR_SSL_WANT_READ;
  }
  return ret;
}

static void mbedtlsDebugPrint(void *ctx, int level, const char *pFile, int line,
                              const char *pStr) {
  Report("\n\rwifi_platform_mbedtls: |%d| %s \n\r", level, pStr);
}

/* ── Init ────────────────────────────────────────────────────── */

int mqtts_tls_init(const char *broker_host, int broker_port) {
  int ret;
  char err[128];
  uint32_t epochTime;
  ip_addr_t serverIp;
  struct sockaddr_in serverAddr;
  // epochTime =  datetime_to_epoch(2026, 4, 20, 9, 21, 28);
  // datetime_SecondsSet( epochTime);
  datetime_printCurTime();
  Report("\nmqtts_tls_init\n");
  Report("\nFree heap before HTTPS: %d bytes\n",
         xPortGetFreeHeapSize()); // FreeRTOS

  mbedtls_ssl_init(&ssl);
  mbedtls_ssl_config_init(&conf);
  mbedtls_x509_crt_init(&ca_cert);
  mbedtls_x509_crt_init(&client_cert);
  mbedtls_pk_init(&client_key);

  psa_crypto_init();

  /* ── TLS config ── */
  ret = mbedtls_ssl_config_defaults(&conf, MBEDTLS_SSL_IS_CLIENT,
                                    MBEDTLS_SSL_TRANSPORT_STREAM,
                                    MBEDTLS_SSL_PRESET_DEFAULT);
  if (ret != 0) {
    mbedtls_strerror(ret, err, sizeof(err));
    UART_PRINT("\n[TLS] config_defaults: %s\n\r", err);
    return ret;
  }
  mbedtls_ssl_conf_max_tls_version(&conf, MBEDTLS_SSL_VERSION_TLS1_2);
  mbedtls_ssl_conf_min_tls_version(&conf, MBEDTLS_SSL_VERSION_TLS1_2);

  // mbedtls debug
  mbedtls_ssl_conf_dbg(&conf, mbedtlsDebugPrint, NULL);
  mbedtls_debug_set_threshold(WIFI_PLATFORM_MBEDTLS_DEBUG_LOG_LEVEL);

  /* ── CA cert (PEM) ── */
  ret = mbedtls_x509_crt_parse(&ca_cert, mosquitto_org_der,
                               mosquitto_org_der_len);
  if (ret != 0) {
    mbedtls_strerror(ret, err, sizeof(err));
    UART_PRINT("\n[TLS] CA cert parse: %s\n\r", err);
    return ret;
  }

  /* ── Client cert (PEM) ── */

  ret = mbedtls_x509_crt_parse(&client_cert, client_der, client_der_len);
  if (ret != 0) {
    mbedtls_strerror(ret, err, sizeof(err));
    UART_PRINT("\n[TLS] client cert parse: %s\n\r", err);
    return ret;
  }

  /* ── Client key (PEM) ── */

  ret = mbedtls_pk_parse_key(&client_key, client_key_der, client_key_der_len,
                             NULL, 0, NULL, NULL);
  if (ret != 0) {
    mbedtls_strerror(ret, err, sizeof(err));
    UART_PRINT("\n[TLS] client key parse: %s\n\r", err);
    return ret;
  }

  mbedtls_ssl_conf_max_frag_len(&conf, MBEDTLS_SSL_MAX_FRAG_LEN_2048);
  mbedtls_ssl_conf_authmode(&conf, MBEDTLS_SSL_VERIFY_REQUIRED);
  mbedtls_ssl_conf_ca_chain(&conf, &ca_cert, NULL);
  mbedtls_ssl_conf_rng(&conf, psa_rng_wrapper, NULL);

  /* ── mTLS: register client cert + key ── */
  ret = mbedtls_ssl_conf_own_cert(&conf, &client_cert, &client_key);
  if (ret != 0) {
    mbedtls_strerror(ret, err, sizeof(err));
    UART_PRINT("\n[TLS] conf_own_cert: %s\n\r", err);
    return ret;
  }

  ret = mbedtls_ssl_setup(&ssl, &conf);
  if (ret != 0) {
    mbedtls_strerror(ret, err, sizeof(err));
    UART_PRINT("\n[TLS] ssl_setup: %s\n\r", err);
    return ret;
  }

  // ret = mbedtls_ssl_set_hostname(&ssl, broker_host);
  ret = mbedtls_ssl_set_hostname(&ssl, "test.mosquitto.org");
  if (ret != 0) {
    mbedtls_strerror(ret, err, sizeof(err));
    UART_PRINT("\n[TLS] set_hostname: %s\n\r", err);
    return ret;
  }

  /* ── TCP connect ── */
  struct addrinfo hints = {.ai_family = AF_INET, .ai_socktype = SOCK_STREAM};
  struct addrinfo *res = NULL;
  char port_str[8];

  snprintf(port_str, sizeof(port_str), "%d", broker_port);

  ret = DNS_IF_gethostbyname(broker_host, &serverIp);
  if (ret != 0) {
    UART_PRINT("[HTTPS] DNS resolution failed: %d\n\r", ret);
    ret = -3;
  }

  // if (lwip_getaddrinfo(broker_host, port_str, &hints, &res) != 0 ||
  //     res == NULL) {
  //     UART_PRINT("\n[TLS] DNS failed for %s\n\r", broker_host);
  //     return -1;
  // }

  sock_fd = lwip_socket(res->ai_family, SOCK_STREAM, 0);
  if (sock_fd < 0) {
    UART_PRINT("\n[TLS] socket() failed\n\r");
    lwip_freeaddrinfo(res);
    return -1;
  }
  /* Set socket timeouts so connect/handshake/read don't block forever */
  {
    struct timeval tv;
    tv.tv_sec = MQTTS_SOCKET_TIMEOUT_MS / 1000;
    tv.tv_usec = (MQTTS_SOCKET_TIMEOUT_MS % 1000) * 1000;
    lwip_setsockopt(sock_fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
    lwip_setsockopt(sock_fd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv));
  }
  memset(&serverAddr, 0, sizeof(serverAddr));
  serverAddr.sin_family = AF_INET;
  serverAddr.sin_port = lwip_htons(broker_port);
  serverAddr.sin_addr.s_addr = ip4_addr_get_u32(ip_2_ip4(&serverIp));

  UART_PRINT("[HTTPS] Connecting to %s:%d...\n\r", broker_host, broker_port);

  ret =
      lwip_connect(sock_fd, (struct sockaddr *)&serverAddr, sizeof(serverAddr));
  // if (lwip_connect(sock_fd, res->ai_addr, res->ai_addrlen) != 0)
  if (ret != 0) {
    UART_PRINT("\n[TLS] TCP connect failed\n\r");
    lwip_freeaddrinfo(res);
    lwip_close(sock_fd);
    sock_fd = -1;
    return -1;
  }
  lwip_freeaddrinfo(res);

  mbedtls_ssl_set_bio(&ssl, &sock_fd, lwip_bio_send, lwip_bio_recv, NULL);
 UART_PRINT("\nmbedtls_platform_set_time...\n\r");
    // mbedtls_time();
    time_t now = time(NULL);
Report("mbedtls_time = %ld\n", now);
  /* ── TLS handshake ── */
  UART_PRINT("\n[TLS] handshake starting...\n\r");
  while ((ret = mbedtls_ssl_handshake(&ssl)) != 0) {
    if (ret != 0) {
      // UART_PRINT("Verify Result: %d\n", mbedtls_ssl_get_verify_result(&ssl));
      uint32_t flags = mbedtls_ssl_get_verify_result(&ssl);

      if (flags != 0) {
        char buf[512];
        mbedtls_x509_crt_verify_info(buf, sizeof(buf), "", flags);
        printf("VERIFY ERROR: %s\n", buf);
      }
    }

    if (ret != MBEDTLS_ERR_SSL_WANT_READ && ret != MBEDTLS_ERR_SSL_WANT_WRITE) {
      mbedtls_strerror(ret, err, sizeof(err));
      UART_PRINT("[TLS] handshake failed: %s (0x%04X)\n\r", err,
                 (unsigned int)-ret);
      mqtts_tls_cleanup();
      return ret;
    }
  }

  UART_PRINT("\n[TLS] handshake OK\n\r");
  return 0;
}

/* ── Cleanup ─────────────────────────────────────────────────── */

void mqtts_tls_cleanup(void) {
  mbedtls_ssl_close_notify(&ssl);

  if (sock_fd >= 0) {
    lwip_close(sock_fd);
    sock_fd = -1;
  }

  mbedtls_ssl_free(&ssl);
  mbedtls_ssl_config_free(&conf);
  mbedtls_x509_crt_free(&ca_cert);
  mbedtls_x509_crt_free(&client_cert);
  mbedtls_pk_free(&client_key);
}
void sntp_set_default_servers(void) {
  // sntpWrapper_store_servers(3,"pool.ntp.org","time.google.com","time.windows.com");
  sntpWrapper_store_servers(3, "129.6.15.28", "216.239.35.0", "13.86.101.172");
}
void sntp_sync_time(void) {

  int32_t ret = sntpWrapper_updateDateTime();

  if (ret < 0) {
    Report("SNTP update failed\n");
  } else {
    Report("SNTP update success\n");
  }
}

/////////////////////////////////////////////////////////////////////////////////////////////////


int32_t ATCmdMqtt_connectCallback(void *arg)
{
    // Report(" Mqtt_connectCallback \n\r");
    char *token;
    char *server;
    uint16_t port;
int32_t ret = 0;

sntp_set_default_servers();

// Sync time
sntp_sync_time();
sleep(2);

   ret =  mqtts_tls_init("test.mosquitto.org", 8884);
    // ret =  mqtts_tls_init("54.36.178.49", 8884);
   Report("MQTT err= %d \n\r",ret);
    if(wlan_connected)
    {
        if(connect_flag != 1)
        {
            char *cmd = (char *)arg;

            token = strtok(cmd, ",");
            if (!token)
                return -1;

            server = token;

            token = strtok(NULL, ",");
            if (!token)
                return -1;

            port = atoi(token);

            // UART_PRINT("MQTT Connect: %s:%d\n\r", server, port);

            /* Create MQTT client */
            gMqttClient = MQTT_IF_clientCreate(NULL, OnMqttEvent);
            ASSERT_GENERAL((uint32_t)gMqttClient);

            MQTTClient_ConnParams connParams = {0};

            connParams.serverAddr = server;
            connParams.port = port;
            connParams.netconnFlags = MQTTCLIENT_NETCONN_URL;
            UART_PRINT("MQTT connecting...\n\r");
            int ret = MQTT_IF_clientConnect(gMqttClient, &connParams);

            if(ret < 0)
            {
                UART_PRINT("MQTT connect failed\n\r");
                return -1;
            }

            
        }
        else 
        {
            UART_PRINT("MQTT already connected\n\r");
        }
    }
    else 
    {
    UART_PRINT("WLAN Not Connected\n\r");
    }
    return 0;
}