#include <Arduino.h>
#include <U8g2lib.h>
#include "HX711.h"

// ---------------- DISPLAY ----------------
// LOW MEMORY DRIVER (_1_ PAGE BUFFER)
U8G2_SSD1306_128X64_NONAME_1_HW_I2C u8g2(U8G2_R0);

// ---------------- PINS ----------------
#define PIN_ENC_CLK   2
#define PIN_ENC_DT    3
#define PIN_ENC_SW    12
#define PIN_INJECTOR  9
#define PIN_PUMP_PWM  10
#define PIN_PUMP_LED  11

// HX711 (ONLY ONE ACTIVE)
#define HX_SCK_PIN    4
#define HX_DT1_PIN    5

// ---------------- UI STATES ----------------
enum UIState {
  UI_STATUS,
  UI_MENU,
  UI_TEST_SETUP,
  UI_STABILIZE,
  UI_EMPTY_TUBES,
  UI_TEST_RUNNING,
  UI_TEST_DONE
};

UIState uiState = UI_STATUS;

// ---------------- MENU NAV ----------------
int menuIndex = 0;   // 0 = Start Test, 1 = Test Setup

// ---------------- TEST SETUP NAV ----------------
int setupIndex = 0;        // 0=PW, 1=Freq, 2=Time, 3=SaveReturn, 4=SaveStart
bool editingValue = false;

// ---------------- ENCODER ----------------
volatile int8_t encDelta = 0;
uint8_t lastEncState = 0;
bool encButtonPrev = HIGH;

void encoderISR() {
  uint8_t a = digitalRead(PIN_ENC_CLK);
  uint8_t b = digitalRead(PIN_ENC_DT);
  uint8_t s = (a << 1) | b;
  uint8_t c = (lastEncState << 2) | s;

  if (c == 0b0001 || c == 0b0111 || c == 0b1110 || c == 0b1000) encDelta++;
  else if (c == 0b0010 || c == 0b0100 || c == 0b1101 || c == 0b1011) encDelta--;

  lastEncState = s;
}

uint8_t readEncoderEvent() {
  uint8_t ev = 0;
  noInterrupts();
  int8_t d = encDelta;
  encDelta = 0;
  interrupts();

  if (d > 0) ev |= 0x01;
  if (d < 0) ev |= 0x02;

  bool sw = digitalRead(PIN_ENC_SW);
  if (sw == LOW && encButtonPrev == HIGH) ev |= 0x04;
  encButtonPrev = sw;

  return ev;
}

// ---------------- PRESSURE ----------------
float readPressurePSI() {
  return 55.0f;  // stable simulation
}

// ---------------- INJECTOR TIMING ----------------
bool injectorRunning = false;
uint32_t cycleStart_us = 0;
uint32_t period_us = 50000;
uint32_t onTime_us = 2000;
bool injectorState = false;

void updateInjectorTiming() {
  if (!injectorRunning) {
    digitalWrite(PIN_INJECTOR, HIGH);
    injectorState = false;
    return;
  }

  uint32_t now = micros();
  uint32_t elapsed = now - cycleStart_us;

  if (elapsed >= period_us) {
    cycleStart_us = now;
    digitalWrite(PIN_INJECTOR, LOW);
    injectorState = true;
    return;
  }

  if (injectorState && elapsed >= onTime_us) {
    digitalWrite(PIN_INJECTOR, HIGH);
    injectorState = false;
  }
}

// ---------------- PUMP CONTROL ----------------
float cfg_targetPressure = 55.0;
float Kp = 1.2;
int pumpDuty = 0;
bool pumpRelayState = false;

void updatePumpControl() {
  if (!pumpRelayState) {
    analogWrite(PIN_PUMP_PWM, 0);
    digitalWrite(PIN_PUMP_LED, HIGH);
    return;
  }

  float p = readPressurePSI();
  float error = cfg_targetPressure - p;

  pumpDuty += Kp * error;
  pumpDuty = constrain(pumpDuty, 0, 100);

  int pwm = map(pumpDuty, 0, 100, 0, 255);
  analogWrite(PIN_PUMP_PWM, pwm);

  digitalWrite(PIN_PUMP_LED, pumpDuty > 0 ? LOW : HIGH);
}

// ---------------- TEST CONFIG ----------------
float testPW_ms = 3.0;
float testFreq_Hz = 10.0;
int testDuration_s = 30;

uint32_t testStartMillis = 0;
uint32_t stabilizeStartMillis = 0;
bool pressureStable = false;

// ---------------- HX711 (ONLY ONE ACTIVE) ----------------
HX711 scale1;

float cal1 = 1000.0f;
float mass1 = 0;

float readScale(HX711 &scale, float cal) {
  if (!scale.is_ready()) return 0.0f;
  long raw = scale.read();
  return (float)raw / cal;
}

void tareScale() {
  scale1.tare();
}

// ---------------- RENDER ----------------
void renderStatus() {
  float p = readPressurePSI();

  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_6x12_tf);

  u8g2.setCursor(0, 20);
  u8g2.print("PSI: ");
  u8g2.print(p);

  u8g2.setCursor(0, 40);
  u8g2.print("CLICK: Menu");

  u8g2.sendBuffer();
}

void renderMenu() {
  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_6x12_tf);

  u8g2.setCursor(0, 20);
  u8g2.print(menuIndex == 0 ? "> Start Test" : "  Start Test");

  u8g2.setCursor(0, 36);
  u8g2.print(menuIndex == 1 ? "> Test Setup" : "  Test Setup");

  u8g2.sendBuffer();
}

void renderTestSetup() {
  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_6x12_tf);

  u8g2.setCursor(0, 10);
  u8g2.print(setupIndex == 0 ? "> PW: " : "  PW: ");
  u8g2.print(testPW_ms, 1);
  u8g2.print(" ms");

  u8g2.setCursor(0, 22);
  u8g2.print(setupIndex == 1 ? "> Freq: " : "  Freq: ");
  u8g2.print(testFreq_Hz, 1);
  u8g2.print(" Hz");

  u8g2.setCursor(0, 34);
  u8g2.print(setupIndex == 2 ? "> Time: " : "  Time: ");
  u8g2.print(testDuration_s);
  u8g2.print(" s");

  u8g2.setCursor(0, 46);
  u8g2.print(setupIndex == 3 ? "> Save & Return" : "  Save & Return");

  u8g2.setCursor(0, 58);
  u8g2.print(setupIndex == 4 ? "> Save & Start" : "  Save & Start");

  u8g2.sendBuffer();
}

void renderStabilizeScreen(float p) {
  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_6x12_tf);

  u8g2.setCursor(0, 20);
  u8g2.print("Stabilizing...");

  u8g2.setCursor(0, 36);
  u8g2.print("PSI: ");
  u8g2.print(p);

  u8g2.sendBuffer();
}

void renderEmptyTubesScreen() {
  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_6x12_tf);

  u8g2.setCursor(0, 20);
  u8g2.print("Empty Tubes");

  u8g2.setCursor(0, 36);
  u8g2.print("Remove fuel");

  u8g2.setCursor(0, 52);
  u8g2.print("CLICK to proceed");

  u8g2.sendBuffer();
}

void renderTestRunningScreen(uint32_t elapsed) {
  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_6x12_tf);

  u8g2.setCursor(0, 20);
  u8g2.print("Test Running");

  u8g2.setCursor(0, 36);
  u8g2.print("Time: ");
  u8g2.print(elapsed / 1000);
  u8g2.print("/");
  u8g2.print(testDuration_s);

  u8g2.sendBuffer();
}

void renderTestResults() {
  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_6x12_tf);

  u8g2.setCursor(0, 20);
  u8g2.print("Test Complete");

  u8g2.setCursor(0, 36);
  u8g2.print("Mass: ");
  u8g2.print(mass1, 1);

  u8g2.sendBuffer();
}

// ---------------- TEST SETUP ----------------
void updateTestSetup() {
  uint8_t ev = readEncoderEvent();

  if (!editingValue) {
    if (ev & 0x01) setupIndex = (setupIndex + 1) % 5;
    if (ev & 0x02) setupIndex = (setupIndex + 4) % 5;

    if (ev & 0x04) {
      if (setupIndex == 3) {
        uiState = UI_MENU;
        return;
      }
      if (setupIndex == 4) {
        uiState = UI_STABILIZE;
        pressureStable = false;
        pumpRelayState = true;
        return;
      }
      editingValue = true;
    }
  } 
  else {
    if (setupIndex == 0) {
      if (ev & 0x01) testPW_ms += 0.1;
      if (ev & 0x02) testPW_ms -= 0.1;
      if (testPW_ms < 0.1) testPW_ms = 0.1;
    }
    else if (setupIndex == 1) {
      if (ev & 0x01) testFreq_Hz += 0.5;
      if (ev & 0x02) testFreq_Hz -= 0.5;
      if (testFreq_Hz < 0.5) testFreq_Hz = 0.5;
    }
    else if (setupIndex == 2) {
      if (ev & 0x01) testDuration_s += 1;
      if (ev & 0x02) testDuration_s -= 1;
      if (testDuration_s < 1) testDuration_s = 1;
    }

    if (ev & 0x04) editingValue = false;
  }

  renderTestSetup();
}

// ---------------- STABILIZATION ----------------
void updateStabilize() {
  float p = readPressurePSI();
  pumpRelayState = true;

  if (abs(p - cfg_targetPressure) < 1.0) {
    if (!pressureStable) {
      pressureStable = true;
      stabilizeStartMillis = millis();
    }

    if (millis() - stabilizeStartMillis > 2000) {
      injectorRunning = false;
      digitalWrite(PIN_INJECTOR, HIGH);
      uiState = UI_EMPTY_TUBES;
      return;
    }
  } else {
    pressureStable = false;
  }

  updatePumpControl();
  updateInjectorTiming();
  renderStabilizeScreen(p);
}

// ---------------- EMPTY TUBES ----------------
void updateEmptyTubes() {
  renderEmptyTubesScreen();

  uint8_t ev = readEncoderEvent();
  if (ev & 0x04) {
    tareScale();
    mass1 = 0;

    period_us = (uint32_t)(1000000.0 / testFreq_Hz);
    onTime_us = (uint32_t)(testPW_ms * 1000.0);

    testStartMillis = millis();
    injectorRunning = true;

    uiState = UI_TEST_RUNNING;
  }
}

// ---------------- TEST RUNNING ----------------
void updateTestRunning() {
  uint32_t elapsed = millis() - testStartMillis;

  mass1 += readScale(scale1, cal1);

  if (elapsed >= testDuration_s * 1000UL) {
    injectorRunning = false;
    pumpRelayState = false;
    uiState = UI_TEST_DONE;
    return;
  }

  updatePumpControl();
  updateInjectorTiming();
  renderTestRunningScreen(elapsed);
}

// ---------------- TEST DONE ----------------
void updateTestDone() {
  updatePumpControl();
  renderTestResults();

  uint8_t ev = readEncoderEvent();
  if (ev & 0x04) uiState = UI_MENU;
}

// ---------------- SETUP ----------------
void setup() {
  digitalWrite(PIN_INJECTOR, HIGH);
  digitalWrite(PIN_PUMP_LED, HIGH);

  pinMode(PIN_INJECTOR, OUTPUT);
  pinMode(PIN_PUMP_LED, OUTPUT);

  pinMode(PIN_ENC_CLK, INPUT_PULLUP);
  pinMode(PIN_ENC_DT, INPUT_PULLUP);
  pinMode(PIN_ENC_SW, INPUT_PULLUP);

  pinMode(PIN_PUMP_PWM, OUTPUT);

  attachInterrupt(digitalPinToInterrupt(PIN_ENC_CLK), encoderISR, CHANGE);
  attachInterrupt(digitalPinToInterrupt(PIN_ENC_DT), encoderISR, CHANGE);

  // HX711 SAFE INIT
  pinMode(HX_DT1_PIN, INPUT_PULLUP);
  pinMode(HX_SCK_PIN, OUTPUT);
  digitalWrite(HX_SCK_PIN, LOW);
  delay(50);

  scale1.begin(HX_DT1_PIN, HX_SCK_PIN);
  tareScale();

  u8g2.begin();
}

// ---------------- LOOP ----------------
void loop() {
  uint8_t ev = readEncoderEvent();

  if (uiState == UI_STATUS) {
    if (ev & 0x04) uiState = UI_MENU;
    renderStatus();
    return;
  }

  if (uiState == UI_MENU) {
    if (ev & 0x01) menuIndex = (menuIndex + 1) % 2;
    if (ev & 0x02) menuIndex = (menuIndex + 1) % 2;

    if (ev & 0x04) {
      if (menuIndex == 0) {
        uiState = UI_STABILIZE;
        pressureStable = false;
        pumpRelayState = true;
      } else {
        uiState = UI_TEST_SETUP;
        setupIndex = 0;
        editingValue = false;
      }
    }

    renderMenu();
    return;
  }

  if (uiState == UI_TEST_SETUP) {
    updateTestSetup();
    return;
  }

  if (uiState == UI_STABILIZE) {
    updateStabilize();
    return;
  }

  if (uiState == UI_EMPTY_TUBES) {
    updateEmptyTubes();
    return;
  }

  if (uiState == UI_TEST_RUNNING) {
    updateTestRunning();
    return;
  }

  if (uiState == UI_TEST_DONE) {
    updateTestDone();
    return;
  }
}