#include <WiFi.h>
#include <PubSubClient.h>
#include <ArduinoJson.h>

/*******/
/* Global settings */
/*******/

const int pwm_period_us = 100;
const char* clientId = "listener-0";
const char* mqttTopic = "Project-3-topic-iot";

/*******/
/* Global structures */
/*******/

volatile bool isRegistered = false;
volatile int deadline = -1;
volatile float pwm_value = 0;

int led = 13;
WiFiClient espClient;
PubSubClient mqttClient(espClient);
hw_timer_t *timer;

/*******/
/* Setup functions */
/*******/

void setupLED(int led) {
  pinMode(led, OUTPUT);
}

void setupTimer(hw_timer_t*& timer, void (*timer_callback)(void), int steps) {
  timer = timerBegin(0, 80, true);
  timerAttachInterrupt(timer, timer_callback, true);
  timerAlarmWrite(timer, steps, true);
  timerAlarmEnable(timer);
}

void setupWiFi(void) {
  static const char* ssid = "Wokwi-GUEST";
  static const char* password = "";

  WiFi.begin(ssid, password);
}

void setupMQTTClient(PubSubClient& mqttClient, void (*mqtt_callback)(char*, byte*, unsigned int)) {
  static const char* mqttServer = "broker.hivemq.com";
  static const int mqttPort = 1883;

  mqttClient.setServer(mqttServer, mqttPort);
  mqttClient.setCallback(mqtt_callback);
}

/***/
/* Utils */
/***/

void awaitConnection(PubSubClient& mqttClient, const char* mqttTopic) {
  if (WiFi.status() != WL_CONNECTED) {
    do {
      Serial.println("Connecting to WiFi...");
      delay(1000);
    } while (WiFi.status() != WL_CONNECTED);

    Serial.println("Connected to WiFi");
  }

  // Loop until connection is completed.
  while (!mqttClient.connected()) {
    Serial.println("Connecting to MQTT broker...");
    if (mqttClient.connect(clientId)) {
      Serial.println("Connected to MQTT broker");
      // Subscribe after connection.
      mqttClient.subscribe(mqttTopic, 1);
    } else {
      Serial.print("Failed, rc=");
      Serial.println(mqttClient.state());
      delay(2000);
    }
  }

  mqttClient.loop();
}

void mqttPublish(PubSubClient& mqttClient, const char* mqttTopic, const char* payload) {
  mqttClient.publish(mqttTopic, payload);
}

/****/
/* Main */
/****/

void IRAM_ATTR timer_callback() {
  bool led_status = digitalRead(led);

  // Tune the led based on the pwm_value variable.
  if (led_status && pwm_value != pwm_period_us) {
    digitalWrite(led, !led_status);
    timerWrite(timer, pwm_value);
  } else if (!led_status && pwm_value != 0) {
    digitalWrite(led, !led_status);
    timerWrite(timer, pwm_period_us - pwm_value);
  }
}

void mqtt_callback(char* topic, byte* payload, unsigned int length) {
  // Convert the received payload to a JSON object
  StaticJsonDocument<200> jsonDoc;
  deserializeJson(jsonDoc, payload, length);

  const char *type = jsonDoc["message_type"];
  Serial.print("Received a message of type: ");
  Serial.println(type);

  // If the message is of type beacon.
  if (strcmp(type, "beacon") == 0) {
    // Check if already registered.
    isRegistered = jsonDoc["slot_allocation"].containsKey(clientId);
    int capDuration = jsonDoc["collision_access_part_duration"];

    if (!isRegistered) {
      // If it is not registered, compute the delay to send a registration message.
      // This rundom delay is chosen between 0 and (something less than) the duration of the CAP part.
      deadline = rand() % ((int)(capDuration * .8)) + millis();

      Serial.print("Client is not registered, send register request of clientId: ");
      Serial.println(clientId);
    }
  }

  // If the message is of type data.
  if (strcmp(type, "data") == 0) {
    const char* to_client_id = jsonDoc["to_client_id"];
    if (strcmp(to_client_id, clientId) == 0) {
      // Tune the pwm _value variable.
      pwm_value = map((float)jsonDoc["humidity"], 0, 100, 0, pwm_period_us);

      Serial.print("Received humidity data: ");
      Serial.println((float)jsonDoc["humidity"]);
    }
  }
}

void setup() {
  Serial.begin(115200);

  srand(time(NULL));
  setupLED(led);
  setupWiFi();
  setupMQTTClient(mqttClient, mqtt_callback);
  setupTimer(timer, timer_callback, pwm_period_us);
}

void loop() {
  awaitConnection(mqttClient, mqttTopic);
  delay(10);

  // Await the delay to pass.
  if (deadline != -1 && millis() > deadline) {
    Serial.println("Timer fired, registration");
    Serial.println(deadline);
    // Reset the delay.
    deadline = -1;

    StaticJsonDocument<200> jsonDoc;
    // Set the values for the JSON object
    jsonDoc["message_type"] = "registration";
    jsonDoc["client_id"] = clientId;
    // Determine the required buffer size for serialization
    size_t bufferSize = measureJson(jsonDoc) + 1;
    // Create a char array as the destination buffer
    char jsonString[bufferSize];

    // Serialize the JSON object to the char array
    serializeJson(jsonDoc, jsonString, bufferSize);

    mqttPublish(mqttClient, mqttTopic, jsonString);
  }
}