#include <Arduino.h>
#include <Wire.h>
#include <SPI.h>
#include <TFT_eSPI.h>
#include <math.h>

// ESP32-2432S028R (2.8 inch CYD) + INA226 red module.
// I2C connector CN1: GND - GPIO22(SCL) - GPIO27(SDA) - 3V3.
// Current path: supply+ -> VIN+ -> VIN- -> load+; all GNDs common.

constexpr uint8_t I2C_SDA = 27;
constexpr uint8_t I2C_SCL = 22;
constexpr uint8_t INA_ADDR = 0x44;
// Calibrated after repairing the open shunt-sense trace with a jumper wire.
// Two-point check with a 100 ohm load: 3.58 V and 7.42 V.
constexpr float SHUNT_OHMS = 0.00807f;
constexpr float BUS_WARN_V = 35.0f;
constexpr float SHUNT_WARN_MV = 78.0f;

constexpr uint8_t TOUCH_IRQ = 36;
constexpr uint8_t TOUCH_MOSI = 32;
constexpr uint8_t TOUCH_MISO = 39;
constexpr uint8_t TOUCH_CLK = 25;
constexpr uint8_t TOUCH_CS = 33;
constexpr uint8_t BOOT_BUTTON = 0;

constexpr uint8_t REG_CONFIG = 0x00;
constexpr uint8_t REG_SHUNT_V = 0x01;
constexpr uint8_t REG_BUS_V = 0x02;
constexpr uint8_t REG_MANUFACTURER = 0xFE;
constexpr uint8_t REG_DIE_ID = 0xFF;

// This CYD clone has red/blue swapped in direct TFT_eSPI colors.
constexpr uint16_t C_BLACK = 0x0000;
constexpr uint16_t C_BG = 0x18C1;
constexpr uint16_t C_PANEL = 0x2983;
constexpr uint16_t C_LINE = 0x4A86;
constexpr uint16_t C_WHITE = 0xFFFF;
constexpr uint16_t C_MUTED = 0xAD75;
constexpr uint16_t C_GREEN = 0x07E0;
constexpr uint16_t C_BLUE = 0xFD20;
constexpr uint16_t C_YELLOW = 0x07FF;
constexpr uint16_t C_RED = 0x001F;

TFT_eSPI tft;
SPIClass touchSPI(VSPI);

struct Reading {
  float busV = 0;
  float shuntMv = 0;
  float currentA = 0;
  float powerW = 0;
};

Reading reading;
bool sensorOnline = false;
bool screenReady = false;
bool touchWasDown = false;
float zeroOffsetMv = 0;
float energyMWh = 0;
float chargeMAh = 0;
float peakW = 0;
float peakA = 0;
uint32_t startedMs = 0;
uint32_t lastMeasureMs = 0;
uint32_t lastDrawMs = 0;
uint32_t lastSerialMs = 0;
uint32_t lastReconnectMs = 0;
String actionMessage;
uint32_t actionUntilMs = 0;

static void setText(uint16_t color, uint8_t size = 1) {
  tft.setTextColor(color, C_BG);
  tft.setTextSize(size);
}

static bool writeRegister(uint8_t reg, uint16_t value) {
  Wire.beginTransmission(INA_ADDR);
  Wire.write(reg);
  Wire.write((uint8_t)(value >> 8));
  Wire.write((uint8_t)value);
  return Wire.endTransmission() == 0;
}

static bool readRegister(uint8_t reg, uint16_t &value) {
  Wire.beginTransmission(INA_ADDR);
  Wire.write(reg);
  if (Wire.endTransmission(false) != 0) return false;
  if (Wire.requestFrom((int)INA_ADDR, 2) != 2) return false;
  value = ((uint16_t)Wire.read() << 8) | Wire.read();
  return true;
}

static bool initIna226() {
  uint16_t manufacturer = 0;
  uint16_t dieId = 0;
  if (!readRegister(REG_MANUFACTURER, manufacturer)) return false;
  if (!readRegister(REG_DIE_ID, dieId)) return false;
  if (manufacturer != 0x5449 || (dieId & 0xFFF0) != 0x2260) {
    Serial.printf("Unexpected ID: MFG=0x%04X DIE=0x%04X\n", manufacturer, dieId);
    return false;
  }
  // AVG=16, bus/shunt conversion=1.1ms, continuous bus+shunt mode.
  if (!writeRegister(REG_CONFIG, 0x0527)) return false;
  Serial.printf("INA226 online at 0x%02X, shunt=%.3f ohm\n", INA_ADDR, SHUNT_OHMS);
  return true;
}

static bool sampleIna226(Reading &out) {
  uint16_t rawBus = 0;
  uint16_t rawShunt = 0;
  if (!readRegister(REG_BUS_V, rawBus)) return false;
  if (!readRegister(REG_SHUNT_V, rawShunt)) return false;
  out.busV = rawBus * 0.00125f;
  out.shuntMv = (int16_t)rawShunt * 0.0025f;
  out.currentA = ((out.shuntMv - zeroOffsetMv) / 1000.0f) / SHUNT_OHMS;
  out.powerW = out.busV * out.currentA;
  return isfinite(out.busV) && isfinite(out.currentA);
}

static String valueText(float value, uint8_t decimals) {
  if (!isfinite(value)) return "--";
  return String(value, (unsigned int)decimals);
}

static String elapsedText() {
  uint32_t seconds = (millis() - startedMs) / 1000;
  uint32_t hours = seconds / 3600;
  uint8_t minutes = (seconds / 60) % 60;
  uint8_t secs = seconds % 60;
  char text[16];
  snprintf(text, sizeof(text), "%02lu:%02u:%02u", (unsigned long)hours, minutes, secs);
  return String(text);
}

static void drawFrame() {
  tft.fillScreen(C_BG);
  tft.setTextDatum(TL_DATUM);
  setText(C_WHITE, 2);
  tft.drawString("INA226 POWER MONITOR", 8, 6, 1);
  tft.drawFastHLine(8, 26, 304, C_LINE);

  tft.fillRoundRect(7, 32, 306, 61, 7, C_PANEL);
  tft.fillRoundRect(7, 99, 150, 52, 7, C_PANEL);
  tft.fillRoundRect(163, 99, 150, 52, 7, C_PANEL);
  tft.fillRoundRect(7, 157, 306, 43, 7, C_PANEL);

  setText(C_MUTED, 1);
  tft.drawString("POWER", 16, 39, 2);
  tft.drawString("VOLTAGE", 16, 105, 2);
  tft.drawString("CURRENT", 172, 105, 2);
  tft.drawString("ENERGY", 15, 162, 1);
  tft.drawString("CHARGE", 91, 162, 1);
  tft.drawString("PEAK W", 167, 162, 1);
  tft.drawString("TIME", 242, 162, 1);

  tft.fillRoundRect(170, 207, 66, 27, 5, C_PANEL);
  tft.fillRoundRect(242, 207, 71, 27, 5, C_PANEL);
  tft.drawRoundRect(170, 207, 66, 27, 5, C_BLUE);
  tft.drawRoundRect(242, 207, 71, 27, 5, C_YELLOW);
  tft.setTextDatum(MC_DATUM);
  setText(C_BLUE, 1);
  tft.drawString("ZERO", 203, 220, 2);
  setText(C_YELLOW, 1);
  tft.drawString("RESET", 277, 220, 2);
  tft.setTextDatum(TL_DATUM);
  screenReady = true;
}

static void clearValueArea(int16_t x, int16_t y, int16_t w, int16_t h, uint16_t color = C_PANEL) {
  tft.fillRect(x, y, w, h, color);
}

static void drawActionButtons() {
  const bool messageActive = millis() < actionUntilMs;
  const bool zeroOk = messageActive && actionMessage == "ZERO OK";
  const bool zeroFail = messageActive && actionMessage == "ZERO FAIL";
  const bool resetOk = messageActive && actionMessage == "RESET OK";

  uint16_t zeroFill = zeroOk ? C_GREEN : (zeroFail ? C_RED : C_PANEL);
  uint16_t zeroEdge = zeroOk ? C_GREEN : (zeroFail ? C_RED : C_BLUE);
  uint16_t zeroText = (zeroOk || zeroFail) ? C_BLACK : C_BLUE;
  uint16_t resetFill = resetOk ? C_GREEN : C_PANEL;
  uint16_t resetEdge = resetOk ? C_GREEN : C_YELLOW;
  uint16_t resetText = resetOk ? C_BLACK : C_YELLOW;

  tft.fillRoundRect(170, 207, 66, 27, 5, zeroFill);
  tft.fillRoundRect(242, 207, 71, 27, 5, resetFill);
  tft.drawRoundRect(170, 207, 66, 27, 5, zeroEdge);
  tft.drawRoundRect(242, 207, 71, 27, 5, resetEdge);
  tft.setTextDatum(MC_DATUM);
  setText(zeroText, 1);
  tft.setTextColor(zeroText, zeroFill);
  tft.drawString("ZERO", 203, 220, 2);
  tft.setTextColor(resetText, resetFill);
  tft.drawString("RESET", 277, 220, 2);
}

static void drawValues() {
  if (!screenReady) drawFrame();

  const bool warning = fabsf(reading.shuntMv) >= SHUNT_WARN_MV || reading.busV >= BUS_WARN_V;
  const uint16_t statusColor = sensorOnline ? (warning ? C_RED : C_GREEN) : C_RED;
  clearValueArea(254, 4, 59, 20, C_BG);
  tft.fillCircle(262, 13, 4, statusColor);
  tft.setTextDatum(TR_DATUM);
  setText(statusColor, 1);
  tft.drawString(sensorOnline ? (warning ? "LIMIT" : "0x44 OK") : "NO 0x44", 312, 7, 1);

  if (!sensorOnline) {
    clearValueArea(12, 52, 296, 35);
    tft.setTextDatum(MC_DATUM);
    setText(C_RED, 2);
    tft.drawString("INA226 NOT FOUND", 160, 70, 2);
    clearValueArea(12, 119, 140, 27);
    clearValueArea(168, 119, 140, 27);
    clearValueArea(12, 174, 296, 22);
  } else {
    clearValueArea(60, 43, 238, 45);
    tft.setTextDatum(MR_DATUM);
    setText(warning ? C_RED : C_GREEN, 1);
    tft.drawString(valueText(reading.powerW, fabsf(reading.powerW) < 10 ? 3 : 2) + " W", 298, 66, 4);

    clearValueArea(12, 120, 140, 27);
    tft.setTextDatum(MR_DATUM);
    setText(C_BLUE, 1);
    tft.drawString(valueText(reading.busV, 3) + " V", 147, 133, 4);

    clearValueArea(168, 120, 140, 27);
    tft.setTextDatum(MR_DATUM);
    setText(C_YELLOW, 1);
    const float absA = fabsf(reading.currentA);
    String current = absA < 1.0f ? valueText(reading.currentA * 1000.0f, 1) + " mA" : valueText(reading.currentA, 3) + " A";
    tft.drawString(current, 303, 133, absA < 1.0f ? 2 : 4);

    clearValueArea(12, 174, 296, 22);
    tft.setTextDatum(MC_DATUM);
    setText(C_WHITE, 1);
    tft.drawString(valueText(energyMWh, 1), 42, 184, 1);
    tft.drawString(valueText(chargeMAh, 1), 118, 184, 1);
    tft.drawString(valueText(peakW, 2), 194, 184, 1);
    tft.drawString(elapsedText(), 275, 184, 1);
  }

  clearValueArea(7, 207, 157, 27, C_BG);
  tft.setTextDatum(TL_DATUM);
  setText(statusColor, 1);
  String footer;
  if (millis() < actionUntilMs) footer = actionMessage;
  else if (!sensorOnline) footer = "CHECK SDA/SCL";
  else if (warning) footer = "OVER RANGE!";
  else if (fabsf(reading.currentA) >= 0.010f) footer = "LOAD ON";
  else footer = "READY NO LOAD";
  tft.drawString(footer, 8, 211, 2);
  drawActionButtons();
}

static void touchInit() {
  touchSPI.begin(TOUCH_CLK, TOUCH_MISO, TOUCH_MOSI, TOUCH_CS);
  pinMode(TOUCH_CS, OUTPUT);
  pinMode(TOUCH_IRQ, INPUT_PULLUP);
  digitalWrite(TOUCH_CS, HIGH);
}

static int16_t bestTwoAverage(int16_t a, int16_t b, int16_t c) {
  const int16_t ab = abs(a - b);
  const int16_t ac = abs(a - c);
  const int16_t bc = abs(b - c);
  if (ab <= ac && ab <= bc) return (a + b) / 2;
  if (ac <= ab && ac <= bc) return (a + c) / 2;
  return (b + c) / 2;
}

static bool readTouch(int &x, int &y) {
  int16_t data[6] = {0, 0, 0, 0, 0, 0};
  touchSPI.beginTransaction(SPISettings(2000000, MSBFIRST, SPI_MODE0));
  digitalWrite(TOUCH_CS, LOW);

  touchSPI.transfer(0xB1);  // Z1
  const int16_t z1 = touchSPI.transfer16(0xC1) >> 3;  // Z2 command follows
  int16_t pressure = z1 + 4095;
  const int16_t z2 = touchSPI.transfer16(0x91) >> 3;  // first X
  pressure -= z2;

  if (pressure >= 300) {
    touchSPI.transfer16(0x91);  // discard first noisy coordinate
    data[0] = touchSPI.transfer16(0xD1) >> 3;
    data[1] = touchSPI.transfer16(0x91) >> 3;
    data[2] = touchSPI.transfer16(0xD1) >> 3;
    data[3] = touchSPI.transfer16(0x91) >> 3;
  }
  data[4] = touchSPI.transfer16(0xD0) >> 3;
  data[5] = touchSPI.transfer16(0x00) >> 3;
  digitalWrite(TOUCH_CS, HIGH);
  touchSPI.endTransaction();

  if (pressure < 300) return false;
  const int16_t rawX = bestTwoAverage(data[0], data[2], data[4]);
  const int16_t rawY = bestTwoAverage(data[1], data[3], data[5]);
  if (rawX < 100 || rawX > 4000 || rawY < 100 || rawY > 4000) return false;

  // Same rotation=1 and calibration used by the verified local CYD project.
  x = constrain(map(rawX, 200, 3700, 0, 320), 0L, 319L);
  y = constrain(map(rawY, 240, 3800, 0, 240), 0L, 239L);
  return true;
}

static void resetTotals() {
  energyMWh = 0;
  chargeMAh = 0;
  peakW = 0;
  peakA = 0;
  startedMs = millis();
  actionMessage = "RESET OK";
  actionUntilMs = millis() + 1800;
  Serial.println("Totals and peaks reset");
  drawValues();
}

static void zeroCurrent() {
  if (!sensorOnline) return;
  float sum = 0;
  uint8_t valid = 0;
  for (uint8_t i = 0; i < 40; ++i) {
    uint16_t raw = 0;
    if (readRegister(REG_SHUNT_V, raw)) {
      sum += (int16_t)raw * 0.0025f;
      ++valid;
    }
    delay(30);
  }
  if (valid) zeroOffsetMv = sum / valid;
  actionMessage = valid ? "ZERO OK" : "ZERO FAIL";
  actionUntilMs = millis() + 1800;
  Serial.printf("Zero offset set to %.4f mV (remove load before ZERO)\n", zeroOffsetMv);
  drawValues();
}

static void handleInput() {
  static int lastTouchX = 0;
  static int lastTouchY = 0;
  int x = 0, y = 0;
  bool down = readTouch(x, y);
  if (down && !touchWasDown) Serial.printf("TOUCH DOWN x=%d y=%d irq=%d\n", x, y, digitalRead(TOUCH_IRQ));
  if (down) {
    lastTouchX = x;
    lastTouchY = y;
  }
  if (!down && touchWasDown) {
    Serial.printf("TOUCH RELEASE x=%d y=%d\n", lastTouchX, lastTouchY);
    if (lastTouchY >= 202 && lastTouchY <= 239 && lastTouchX >= 155 && lastTouchX < 242) zeroCurrent();
    if (lastTouchY >= 202 && lastTouchY <= 239 && lastTouchX >= 242) resetTotals();
  }
  touchWasDown = down;

  static bool bootWasDown = false;
  static uint32_t bootDownMs = 0;
  bool bootDown = digitalRead(BOOT_BUTTON) == LOW;
  if (bootDown && !bootWasDown) bootDownMs = millis();
  if (!bootDown && bootWasDown) {
    if (millis() - bootDownMs >= 1500) resetTotals();
    else zeroCurrent();
  }
  bootWasDown = bootDown;
}

void setup() {
  Serial.begin(115200);
  pinMode(TFT_BL, OUTPUT);
  digitalWrite(TFT_BL, HIGH);
  pinMode(BOOT_BUTTON, INPUT_PULLUP);

  tft.init();
  tft.setRotation(1);
  tft.setTextWrap(false);
  drawFrame();

  touchInit();

  Wire.begin(I2C_SDA, I2C_SCL);
  Wire.setClock(100000);
  Wire.setTimeOut(50);
  sensorOnline = initIna226();
  startedMs = lastMeasureMs = millis();
  drawValues();
  Serial.println("ms,bus_V,shunt_mV,current_A,power_W,charge_mAh,energy_mWh,peak_A,peak_W");
}

void loop() {
  handleInput();
  const uint32_t now = millis();

  if (!sensorOnline && now - lastReconnectMs >= 1000) {
    lastReconnectMs = now;
    sensorOnline = initIna226();
    if (sensorOnline) lastMeasureMs = now;
  }

  if (sensorOnline && now - lastMeasureMs >= 100) {
    const float dtHours = (now - lastMeasureMs) / 3600000.0f;
    lastMeasureMs = now;
    Reading next;
    if (sampleIna226(next)) {
      reading = next;
      if (dtHours <= (2.0f / 3600.0f)) {
        chargeMAh += reading.currentA * 1000.0f * dtHours;
        energyMWh += reading.powerW * 1000.0f * dtHours;
      }
      peakA = max(peakA, fabsf(reading.currentA));
      peakW = max(peakW, fabsf(reading.powerW));
    } else {
      sensorOnline = false;
      Serial.println("INA226 read failed; reconnecting...");
    }
  }

  if (now - lastDrawMs >= 500) {
    lastDrawMs = now;
    drawValues();
  }

  if (sensorOnline && now - lastSerialMs >= 1000) {
    lastSerialMs = now;
    Serial.printf("%lu,%.4f,%.4f,%.5f,%.5f,%.3f,%.3f,%.5f,%.5f\n",
                  (unsigned long)now, reading.busV, reading.shuntMv,
                  reading.currentA, reading.powerW, chargeMAh, energyMWh,
                  peakA, peakW);
  }

  delay(5);
}