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

// ================= NETWORK =================
const char* ssid = "Wokwi-GUEST";
const char* password = "";
const char* mqtt_server = "broker.emqx.io";

const char* topic_req = "wokwibank/atm999_secure/request";
const char* topic_res = "wokwibank/atm999_secure/response";

WiFiClient espClient;
PubSubClient mqtt(espClient);

// ================= LCD & KEYPAD =================
LiquidCrystal_I2C lcd(0x27, 20, 4);

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] = {19, 18, 5, 17};
byte colPins[COLS] = {16, 4, 2, 15};
Keypad keypad = Keypad(makeKeymap(keys), rowPins, colPins, ROWS, COLS);

// ================= STATES & DATA =================
enum ATMState { ENTER_ACC, ENTER_PIN, MENU, ENTER_AMOUNT, PROCESSING, SHOW_RESULT };
ATMState currentState = ENTER_ACC;

String accBuffer = "";
String pinBuffer = "";
String amtBuffer = "";
char transactionType = ' ';

// ================= TIMERS =================
unsigned long requestTime = 0;
unsigned long resultDisplayTime = 0;
bool waitingResponse = false;

// ================= LCD UI =================
void updateLCD() {
  lcd.clear();
  switch (currentState) {
    case ENTER_ACC:
      lcd.setCursor(0, 0); lcd.print("Welcome to WokwiBank");
      lcd.setCursor(0, 1); lcd.print("Enter Acc No:");
      lcd.setCursor(0, 2); lcd.print(accBuffer);
      lcd.setCursor(0, 3); lcd.print("# to confirm");
      break;
    case ENTER_PIN:
      lcd.setCursor(0, 0); lcd.print("Acc: " + accBuffer);
      lcd.setCursor(0, 1); lcd.print("Enter PIN:");
      lcd.setCursor(0, 2);
      for (int i = 0; i < pinBuffer.length(); i++) lcd.print('*');
      lcd.setCursor(0, 3); lcd.print("# to confirm");
      break;
    case MENU:
      lcd.setCursor(0, 0); lcd.print("1.Withdraw");
      lcd.setCursor(0, 1); lcd.print("2.Deposit");
      lcd.setCursor(0, 2); lcd.print("3.Balance");
      lcd.setCursor(0, 3); lcd.print("A.Cancel");
      break;
    case ENTER_AMOUNT:
      lcd.setCursor(0, 0); lcd.print(transactionType == 'W' ? "Withdraw Amt:" : "Deposit Amt:");
      lcd.setCursor(0, 1); lcd.print("$" + amtBuffer);
      lcd.setCursor(0, 3); lcd.print("# to confirm");
      break;
    case PROCESSING:
      lcd.setCursor(0, 1); lcd.print("Processing...");
      lcd.setCursor(0, 2); lcd.print("Please wait...");
      break;
    case SHOW_RESULT:
      break;
  }
}

void resetATM() {
  accBuffer = "";
  pinBuffer = "";
  amtBuffer = "";
  transactionType = ' ';
  waitingResponse = false;
  currentState = ENTER_ACC;
  updateLCD();
}

// ================= MQTT CALLBACK =================
void callback(char* topic, byte* payload, unsigned int length) {
  if (String(topic) != topic_res) return;

  StaticJsonDocument<256> res;
  DeserializationError error = deserializeJson(res, payload, length);

  if (error) {
    Serial.print("JSON parse failed: ");
    Serial.println(error.f_str());
    return;
  }

  waitingResponse = false;
  String status = res["status"];
  String msg = res["msg"];
  int balance = res["balance"];

  currentState = SHOW_RESULT;
  resultDisplayTime = millis(); 

  lcd.clear();
  lcd.setCursor(0, 0); lcd.print(status == "success" ? "TXN SUCCESS" : "TXN FAILED");
  lcd.setCursor(0, 1); lcd.print(msg);
  lcd.setCursor(0, 2); lcd.print("Bal: $" + String(balance));
  lcd.setCursor(0, 3); lcd.print("Returning...");
}

// ================= MQTT CONNECT =================
void reconnect() {
  while (!mqtt.connected()) {
    lcd.clear();
    lcd.setCursor(0, 1); lcd.print("Connecting MQTT...");
    String clientId = "ATM_Client_" + String(random(0xffffff), HEX);
    if (mqtt.connect(clientId.c_str())) {
      mqtt.subscribe(topic_res);
      updateLCD();
    } else {
      delay(1000);
    }
  }
}

// ================= SETUP =================
void setup() {
  Serial.begin(115200);
  lcd.init();
  lcd.backlight();

  WiFi.mode(WIFI_STA);
  WiFi.begin(ssid, password);
  lcd.setCursor(0, 1); lcd.print("Connecting WiFi...");
  
  while (WiFi.status() != WL_CONNECTED) { delay(500); }

  mqtt.setServer(mqtt_server, 1883);
  mqtt.setCallback(callback);
  mqtt.setBufferSize(512);
  
  randomSeed(analogRead(34) + micros());

  reconnect();
  resetATM();
}

// ================= SEND REQUEST =================
// Helper function to safely build JSON 
void sendTransactionRequest(int amount) {
  StaticJsonDocument<200> doc;
  doc["acc"] = accBuffer;
  doc["pin"] = pinBuffer;
  doc["type"] = String(transactionType);
  doc["amount"] = amount;

  String payload;
  serializeJson(doc, payload);

  mqtt.publish(topic_req, payload.c_str());
  requestTime = millis();
  waitingResponse = true;
}

// ================= LOOP =================
void loop() {
  if (!mqtt.connected()) reconnect();
  mqtt.loop();

  if (currentState == SHOW_RESULT) {
    if (millis() - resultDisplayTime > 4000) {
      resetATM();
    }
    return; 
  }

  if (waitingResponse && (millis() - requestTime > 5000)) {
    lcd.clear();
    lcd.setCursor(0, 1); lcd.print("Server Timeout!");
    waitingResponse = false;
    currentState = SHOW_RESULT; 
    resultDisplayTime = millis(); 
    return;
  }

  char key = keypad.getKey();
  if (!key) return;

  if (key == 'A') { resetATM(); return; }

  switch (currentState) {
    case ENTER_ACC:
      if (isdigit(key) && accBuffer.length() < 4) { accBuffer += key; updateLCD(); }
      else if (key == '#' && accBuffer.length() > 0) { currentState = ENTER_PIN; updateLCD(); }
      break;

    case ENTER_PIN:
      if (isdigit(key) && pinBuffer.length() < 4) { pinBuffer += key; updateLCD(); }
      else if (key == '#' && pinBuffer.length() > 0) { currentState = MENU; updateLCD(); }
      break;

    case MENU:
      if (key == '1') { transactionType = 'W'; currentState = ENTER_AMOUNT; updateLCD(); }
      else if (key == '2') { transactionType = 'D'; currentState = ENTER_AMOUNT; updateLCD(); }
      else if (key == '3') {
        transactionType = 'B';
        currentState = PROCESSING;
        updateLCD();
        sendTransactionRequest(0);
      }
      break;

    case ENTER_AMOUNT:
      if (isdigit(key) && amtBuffer.length() < 5) { amtBuffer += key; updateLCD(); }
      else if (key == '#' && amtBuffer.length() > 0) {
        currentState = PROCESSING;
        updateLCD();
        sendTransactionRequest(amtBuffer.toInt());
      }
      break;
      
    default: break;
  }
}