#include <WiFi.h>
#include <PubSubClient.h>
#include <LiquidCrystal_I2C.h>
#include <Keypad.h>

// WiFi
const char* ssid = "Wokwi-GUEST";
const char* password = "";

// MQTT
const char* mqtt_server = "broker.hivemq.com";
const char* topic = "voting/topic";

WiFiClient espClient;
PubSubClient client(espClient);

// LCD
LiquidCrystal_I2C lcd(0x27, 16, 2);

// Keypad
const byte ROWS = 4;
const byte COLS = 4;

char keys[ROWS][COLS] = {
  {'1','2','3','A'},
  {'4','5','6','B'},
  {'7','8','9','C'},
  {'*','0','#','D'}
};

byte rowPins[ROWS] = {13, 12, 14, 27};
byte colPins[COLS] = {26, 25, 33, 32};

Keypad keypad = Keypad(makeKeymap(keys), rowPins, colPins, ROWS, COLS);

// Candidates
struct Candidate {
  String name;
  int votes;
};

Candidate candidates[] = {
  {"ABC", 0},
  {"DEF", 0},
  {"GHI", 0},
  {"JKL", 2},
  {"MNO", 0}
};

unsigned long lastReconnectAttempt = 0;

// UI
void showMainScreen() {
  lcd.clear();
  lcd.setCursor(0,0);
  lcd.print("Waiting Votes");
  lcd.setCursor(0,1);
  lcd.print("Press 1 Result");
}

// WiFi
void setup_wifi() {
  Serial.print("WiFi...");
  WiFi.begin(ssid, password);
  while (WiFi.status() != WL_CONNECTED) {
    delay(300);
    Serial.print(".");
  }
  Serial.println("Connected");
}

// MQTT reconnect
boolean reconnect() {
  Serial.print("MQTT connecting...");

  if (client.connect("ESP32Server")) {
    Serial.println("Connected");

    if (client.subscribe(topic)) {
      Serial.println("Subscribed OK");
    } else {
      Serial.println("Subscribe FAILED");
    }

    showMainScreen();
    return true;
  } else {
    Serial.print("Failed, state=");
    Serial.println(client.state());
    return false;
  }
}

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

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

  Serial.print("Received: ");
  Serial.println(msg);

  int index = msg[0] - '1';

  if (index >= 0 && index < 5) {
    candidates[index].votes++;

    lcd.clear();
    lcd.print("Vote Received");

    delay(500);

    showMainScreen();
  }
}

// Show results
void showResults() {
  Serial.println("Showing Results");

  int maxVotes = -1;
  String winner = "";

  for (int i = 0; i < 5; i++) {

    lcd.clear();
    lcd.print(candidates[i].name);

    lcd.setCursor(0,1);
    lcd.print("Votes:");
    lcd.print(candidates[i].votes);

    delay(700);

    if (candidates[i].votes > maxVotes) {
      maxVotes = candidates[i].votes;
      winner = candidates[i].name;
    }
  }

  lcd.clear();
  lcd.print("Winner:");
  lcd.setCursor(0,1);
  lcd.print(winner);

  delay(1200);

  showMainScreen();
}

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

  lcd.init();
  lcd.backlight();

  setup_wifi();

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

  showMainScreen();
}

void loop() {

  // MQTT
  if (!client.connected()) {
    unsigned long now = millis();
    if (now - lastReconnectAttempt > 3000) {
      lastReconnectAttempt = now;
      reconnect();
    }
  } else {
    client.loop();
  }

  // Keypad
  char key = keypad.getKey();

  if (key) {
    Serial.print("Key pressed: ");
    Serial.println(key);
  }

  if (key == '1') {
    showResults();
  }
}