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

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

const char* ssid = "Wokwi-GUEST";
const char* password = "";
const char* mqtt_server = "broker.hivemq.com";

WiFiClient espClient;
PubSubClient client(espClient);

StaticJsonDocument<4096> db;

void show(String msg) {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0,0);
  display.println(msg);
  display.display();
}

void initDB() {
  JsonArray candidates = db.createNestedArray("candidates");

  for (int i = 1; i <= 5; i++) {
    JsonObject c = candidates.createNestedObject();
    c["id"] = i;
    c["name"] = "Candidate" + String(i);
    c["party"] = "Party" + String(i);
    c["constituency"] = "Chennai";
    c["votes"] = 0;
  }

  JsonArray voters = db.createNestedArray("voters");

  for (int i = 1; i <= 25; i++) {
    JsonObject v = voters.createNestedObject();
    v["id"] = i;
    v["constituency"] = "Chennai";
    v["voted"] = false;
  }
}

void processVote(int voter_id, int candidate_id) {

  JsonArray voters = db["voters"];
  JsonArray candidates = db["candidates"];

  for (JsonObject v : voters) {

    if (v["id"] == voter_id) {

      if (v["voted"] == true) {
        show("Already Voted");
        return;
      }

      v["voted"] = true;

      for (JsonObject c : candidates) {
        if (c["id"] == candidate_id) {
          c["votes"] = c["votes"].as<int>() + 1;
          show("Vote Counted");
          return;
        }
      }
    }
  }

  show("Invalid Voter");
}

void showResults() {
  JsonArray candidates = db["candidates"];

  display.clearDisplay();
  int y = 0;

  for (JsonObject c : candidates) {
    display.setCursor(0, y);
    display.print(c["name"].as<String>());
    display.print(":");
    display.print(c["votes"].as<int>());
    y += 10;
  }

  display.display();
}

void callback(char* topic, byte* payload, unsigned int length) {

  String msg;
  for (int i = 0; i < length; i++) msg += (char)payload[i];

  Serial.println(msg);

  if (String(topic) == "evm/vote/data") {

    StaticJsonDocument<200> doc;
    deserializeJson(doc, msg);

    int voter_id = doc["voter_id"];
    int candidate_id = doc["candidate_id"];

    processVote(voter_id, candidate_id);
  }

  if (String(topic) == "evm/vote/end") {
    showResults();
  }
}

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

  WiFi.begin(ssid, password);
  while (WiFi.status() != WL_CONNECTED);

  client.setServer(mqtt_server, 1883);
  client.setCallback(callback);

  display.begin(SSD1306_SWITCHCAPVCC, 0x3C);

  initDB();
  show("Server Ready");
}

void reconnect() {
  while (!client.connected()) {
    client.connect("EVM_Server");
    client.subscribe("evm/vote/data");
    client.subscribe("evm/vote/end");
  }
}

void loop() {
  if (!client.connected()) reconnect();
  client.loop();
}