#include "mqtt.h"

#include "lwip/altcp.h"
#include "lwip/altcp_tls.h"
#include "lwip/apps/mqtt_priv.h"
#include "lwip/dns.h"
#include "pico/cyw43_arch.h"
#include <string.h>

static void ensaht_mqtt_default_connect_cb(mqtt_client_t *client, void *arg,
                                           mqtt_connection_status_t status) {
  LWIP_UNUSED_ARG(client);
  LWIP_UNUSED_ARG(arg);
  LWIP_UNUSED_ARG(status);
}

static err_t ensaht_mqtt_start(ensaht_mqtt_t *mqtt) {
  err_t err;

  if (mqtt->client == NULL) {
    mqtt->client = mqtt_client_new();
    if (mqtt->client == NULL) {
      return ERR_MEM;
    }
  }

  err = mqtt_client_connect(mqtt->client, &mqtt->server_addr, mqtt->port,
                            mqtt->connect_cb, mqtt->connect_arg,
                            &mqtt->client_info);
#if LWIP_ALTCP && LWIP_ALTCP_TLS
  if (err == ERR_OK && mqtt->client_info.tls_config != NULL) {
    mbedtls_ssl_set_hostname(altcp_tls_context(mqtt->client->conn),
                             mqtt->hostname);
  }
#endif

  return err;
}

static void ensaht_mqtt_dns_found(const char *hostname, const ip_addr_t *ipaddr,
                                  void *arg) {
  ensaht_mqtt_t *mqtt = (ensaht_mqtt_t *)arg;
  LWIP_UNUSED_ARG(hostname);

  if (ipaddr == NULL) {
    if (mqtt->connect_cb != NULL) {
      mqtt->connect_cb(mqtt->client, mqtt->connect_arg, MQTT_CONNECT_TIMEOUT);
    }
    return;
  }

  mqtt->server_addr = *ipaddr;
  cyw43_arch_lwip_begin();
  err_t err = ensaht_mqtt_start(mqtt);
  cyw43_arch_lwip_end();

  if (err != ERR_OK && mqtt->connect_cb != NULL) {
    mqtt->connect_cb(mqtt->client, mqtt->connect_arg,
                     MQTT_CONNECT_DISCONNECTED);
  }
}

void ensaht_mqtt_init(ensaht_mqtt_t *mqtt) { memset(mqtt, 0, sizeof(*mqtt)); }

err_t ensaht_mqtt_connect(ensaht_mqtt_t *mqtt, const char *hostname, u16_t port,
                          const struct mqtt_connect_client_info_t *client_info,
                          mqtt_connection_cb_t cb, void *arg) {
  err_t err;

  if (mqtt == NULL || hostname == NULL || client_info == NULL ||
      client_info->client_id == NULL) {
    return ERR_ARG;
  }

  mqtt->hostname = hostname;
  mqtt->port = port != 0 ? port : MQTT_PORT;
  mqtt->client_info = *client_info;
  mqtt->connect_cb = cb != NULL ? cb : ensaht_mqtt_default_connect_cb;
  mqtt->connect_arg = arg;

  cyw43_arch_lwip_begin();
  err = dns_gethostbyname(hostname, &mqtt->server_addr, ensaht_mqtt_dns_found,
                          mqtt);
  if (err == ERR_OK) {
    err = ensaht_mqtt_start(mqtt);
  }
  cyw43_arch_lwip_end();

  return err == ERR_INPROGRESS ? ERR_OK : err;
}

void ensaht_mqtt_disconnect(ensaht_mqtt_t *mqtt) {
  if (mqtt == NULL || mqtt->client == NULL) {
    return;
  }

  cyw43_arch_lwip_begin();
  mqtt_disconnect(mqtt->client);
  cyw43_arch_lwip_end();
}

err_t ensaht_mqtt_publish(mqtt_client_t *client, const char *topic,
                          const void *payload, u16_t payload_length,
                          mqtt_request_cb_t cb, void *arg) {
  return mqtt_publish(client, topic, payload, payload_length, 1, 0, cb, arg);
}

err_t ensaht_mqtt_publish_client(ensaht_mqtt_t *mqtt, const char *topic,
                                 const void *payload, u16_t payload_length,
                                 mqtt_request_cb_t cb, void *arg) {
  err_t err;

  if (mqtt == NULL || mqtt->client == NULL) {
    return ERR_CONN;
  }

  cyw43_arch_lwip_begin();
  err = ensaht_mqtt_publish(mqtt->client, topic, payload, payload_length, cb,
                            arg);
  cyw43_arch_lwip_end();

  return err;
}

err_t ensaht_mqtt_pingreq(mqtt_client_t *client) {
  static const u8_t pingreq[] = {0xC0, 0x00};
  err_t err;

  LWIP_ASSERT_CORE_LOCKED();
  if (client == NULL || client->conn == NULL ||
      !mqtt_client_is_connected(client)) {
    return ERR_CONN;
  }

  err =
      altcp_write(client->conn, pingreq, sizeof(pingreq), TCP_WRITE_FLAG_COPY);
  if (err == ERR_OK) {
    client->cyclic_tick = 0;
    err = altcp_output(client->conn);
  }

  return err;
}

err_t ensaht_mqtt_pingreq_client(ensaht_mqtt_t *mqtt) {
  err_t err;

  if (mqtt == NULL || mqtt->client == NULL) {
    return ERR_CONN;
  }

  cyw43_arch_lwip_begin();
  err = ensaht_mqtt_pingreq(mqtt->client);
  cyw43_arch_lwip_end();

  return err;
}