/* LVGL + AXP2101 Power Monitor
 * I2C: SDA=8, SCL=7
 * Display: ST7789 240x240 (1.54-inch)
 * Shows: AXP2101 Battery/VBUS/Charge Status
 */

#include <lvgl.h>
#include <Arduino_GFX_Library.h>
#include <Adafruit_MCP23X17.h>
// AXP2101 Library
#define XPOWERS_CHIP_AXP2101
#include <Wire.h>
#include <Arduino.h>
#include "XPowersLib.h"

// Pin definitions
#define GFX_BL 6
#define SPI_MISO 11
#define SPI_MOSI 13
#define SPI_SCLK 12
#define LCD_CS 10
#define LCD_DC 16
#define LCD_RST 10
#define LCD_HOR_RES 240
#define LCD_VER_RES 240
#define I2C_SDA_PIN 8    // AXP2101 I2C SDA
#define I2C_SCL_PIN 7    // AXP2101 I2C SCL
#define PMU_IRQ_PIN -1   // AXP2101 IRQ (unused)

// Global objects
Adafruit_MCP23X17 mcp;
Arduino_DataBus *bus = new Arduino_ESP32SPI(LCD_DC, LCD_CS, SPI_SCLK, SPI_MOSI, SPI_MISO);
Arduino_GFX *gfx = new Arduino_ST7789(bus, LCD_RST, 0, true, LCD_HOR_RES, LCD_VER_RES);
XPowersPMU power; // AXP2101 PMU object

// LVGL display buffer
uint32_t screenWidth, screenHeight, bufSize;
lv_disp_draw_buf_t draw_buf;
lv_color_t *disp_draw_buf1;
lv_color_t *disp_draw_buf2;
lv_disp_drv_t disp_drv;

// UI elements (only AXP2101 related)
lv_obj_t *label_title;
lv_obj_t *label_batt_vol;  // Battery Voltage
lv_obj_t *label_vbus_vol;  // VBUS Voltage
lv_obj_t *label_batt_per;  // Battery Percentage
lv_obj_t *label_charge;    // Charging Status
lv_obj_t *label_sys_vol;   // System Voltage
lv_obj_t *label_status;    // Global Status
lv_timer_t *axp_timer = NULL;

/* -------------------------- AXP2101 Core Functions -------------------------- */
// Initialize AXP2101 PMU
bool init_axp2101() {
  Serial.println("Initializing AXP2101...");
  bool result = power.begin(Wire, AXP2101_SLAVE_ADDRESS, I2C_SDA_PIN, I2C_SCL_PIN);
  if (result == false) {
    Serial.println("AXP2101 initialization failed!");
    return false;
  }
  Serial.printf("AXP2101 Chip ID: 0x%x\n", power.getChipID());

  // Core AXP2101 Configuration
  power.setVbusVoltageLimit(XPOWERS_AXP2101_VBUS_VOL_LIM_4V36);
  power.setVbusCurrentLimit(XPOWERS_AXP2101_VBUS_CUR_LIM_1500MA);
  power.setSysPowerDownVoltage(2600); // VSYS shutdown voltage 2600mV

  // Enable key voltage outputs
  power.enableDC2(); power.enableDC3(); power.enableDC4(); power.enableDC5();
  power.enableALDO1(); power.enableALDO2(); power.enableALDO3(); power.enableALDO4();
  power.enableBLDO1(); power.enableBLDO2(); power.enableCPUSLDO(); power.enableDLDO1(); power.enableDLDO2();

  // Set voltage values (match original config)
  power.setDC1Voltage(3300); power.setDC2Voltage(1000); power.setDC3Voltage(3300);
  power.setDC4Voltage(1000); power.setDC5Voltage(3300);
  power.setALDO1Voltage(3300); power.setALDO2Voltage(3300); power.setALDO3Voltage(3300); power.setALDO4Voltage(3300);
  power.setBLDO1Voltage(1500); power.setBLDO2Voltage(2800); power.setCPUSLDOVoltage(1000);
  power.setDLDO1Voltage(3300); power.setDLDO2Voltage(3300);

  // Enable ADC measurements
  power.disableTSPinMeasure(); // Disable TS pin if no battery temp sensor
  power.enableBattDetection();
  power.enableVbusVoltageMeasure();
  power.enableBattVoltageMeasure();
  power.enableSystemVoltageMeasure();

  // Charging configuration
  power.setPrechargeCurr(XPOWERS_AXP2101_PRECHARGE_50MA);
  power.setChargerConstantCurr(XPOWERS_AXP2101_CHG_CUR_200MA);
  power.setChargerTerminationCurr(XPOWERS_AXP2101_CHG_ITERM_25MA);
  power.setChargeTargetVoltage(XPOWERS_AXP2101_CHG_VOL_4V1);

  // Interrupt configuration (simplified)
  power.disableIRQ(XPOWERS_AXP2101_ALL_IRQ);
  power.clearIrqStatus();

  Serial.println("AXP2101 initialized successfully");
  return true;
}

// Read AXP2101 key data for display
void read_axp2101_data(
  uint16_t &batt_vol, 
  uint16_t &vbus_vol, 
  uint16_t &sys_vol, 
  uint8_t &batt_per, 
  bool &is_charging
) {
  batt_vol = power.getBattVoltage();       // Battery voltage (mV)
  vbus_vol = power.getVbusVoltage();       // VBUS voltage (mV)
  sys_vol = power.getSystemVoltage();      // System voltage (mV)
  batt_per = power.getBatteryPercent();    // Battery percentage (%)
  is_charging = power.isCharging();        // Charging status (true/false)
}

/* -------------------------- LVGL Display Functions -------------------------- */
// LVGL display flushing callback
void my_disp_flush(lv_disp_drv_t *disp_drv, const lv_area_t *area, lv_color_t *color_p) {
  uint32_t w = (area->x2 - area->x1 + 1);
  uint32_t h = (area->y2 - area->y1 + 1);

#if (LV_COLOR_16_SWAP != 0)
  gfx->draw16bitBeRGBBitmap(area->x1, area->y1, (uint16_t *)&color_p->full, w, h);
#else
  gfx->draw16bitRGBBitmap(area->x1, area->y1, (uint16_t *)&color_p->full, w, h);
#endif

  lv_disp_flush_ready(disp_drv);
}

// LCD reset function
void lcd_reset(void) {
  mcp.digitalWrite(1, HIGH);
  delay(10);
  mcp.digitalWrite(1, LOW);
  delay(10);
  mcp.digitalWrite(1, HIGH);
  delay(50);  // Reduced delay for faster screen clearing
}

/* -------------------------- Sensor Initialization -------------------------- */
// Initialize I2C & AXP2101
bool init_sensors() {
  // Initialize I2C with pull-up
  pinMode(I2C_SDA_PIN, INPUT_PULLUP);
  pinMode(I2C_SCL_PIN, INPUT_PULLUP);
  delay(200);

  Wire.end();
  Wire.begin(I2C_SDA_PIN, I2C_SCL_PIN);
  Wire.setClock(100000); // AXP2101 supports 100kHz I2C clock
  Wire.setTimeout(1000000);
  delay(300);
  Serial.println("I2C initialized (100kHz)");

  // Initialize AXP2101
  if (!init_axp2101()) {
    return false;
  }

  return true;
}

/* -------------------------- UI Helper Functions -------------------------- */
// Convert integer to string (for voltage/percentage)
void int_to_str(char *buf, uint32_t value) {
  sprintf(buf, "%d", value);
}

/* -------------------------- AXP2101 Data Update Callback -------------------------- */
static void axp_data_callback(lv_timer_t *timer) {
  char buf[32];
  bool is_axp_ok = true;

  // Read AXP2101 data
  uint16_t batt_vol = 0, vbus_vol = 0, sys_vol = 0;
  uint8_t batt_per = 0;
  bool is_charging = false;
  read_axp2101_data(batt_vol, vbus_vol, sys_vol, batt_per, is_charging);

  // Validate data
  if (batt_vol == 0 && vbus_vol == 0) {
    is_axp_ok = false;
  }

  // Update UI labels
  // Battery Voltage
  if (is_axp_ok) {
    sprintf(buf, "Batt Vol: %d mV", batt_vol);
    lv_label_set_text(label_batt_vol, buf);
  } else {
    lv_label_set_text(label_batt_vol, "Batt Vol: -- mV");
  }

  // VBUS Voltage
  if (is_axp_ok) {
    sprintf(buf, "VBUS Vol: %d mV", vbus_vol);
    lv_label_set_text(label_vbus_vol, buf);
  } else {
    lv_label_set_text(label_vbus_vol, "VBUS Vol: -- mV");
  }

  // System Voltage
  if (is_axp_ok) {
    sprintf(buf, "Sys Vol:  %d mV", sys_vol);
    lv_label_set_text(label_sys_vol, buf);
  } else {
    lv_label_set_text(label_sys_vol, "Sys Vol:  -- mV");
  }

  // Battery Percentage
  if (is_axp_ok) {
    sprintf(buf, "Batt%%:    %d %%", batt_per);
    lv_label_set_text(label_batt_per, buf);
  } else {
    lv_label_set_text(label_batt_per, "Batt%%:    -- %%");
  }

  // Charging Status
  if (is_charging && is_axp_ok) {
    lv_label_set_text(label_charge, "Charging: YES");
    lv_obj_set_style_text_color(label_charge, lv_color_hex(0x00FF00), 0); // Green
  } else if (is_axp_ok) {
    lv_label_set_text(label_charge, "Charging: NO");
    lv_obj_set_style_text_color(label_charge, lv_color_hex(0xAAAAAA), 0); // Gray
  } else {
    lv_label_set_text(label_charge, "Charging: --");
  }

  // Global Status
  if (is_axp_ok) {
    lv_label_set_text(label_status, "Status: Normal");
    lv_obj_set_style_text_color(label_status, lv_color_hex(0x00FF00), 0); // Green
  } else {
    lv_label_set_text(label_status, "Status: Error");
    lv_obj_set_style_text_color(label_status, lv_color_hex(0xFF0000), 0); // Red
  }

  // Serial output (only AXP2101 data)
  Serial.println("\n==================== AXP2101 Data ====================");
  if (is_axp_ok) {
    Serial.printf("Battery Voltage: %d mV\n", batt_vol);
    Serial.printf("VBUS Voltage:    %d mV\n", vbus_vol);
    Serial.printf("System Voltage:  %d mV\n", sys_vol);
    Serial.printf("Battery Percent: %d %%\n", batt_per);
    Serial.printf("Charging:        %s\n", is_charging ? "YES" : "NO");
  } else {
    Serial.println("Status: No Valid Data (AXP2101 not responding)");
  }
  Serial.println("======================================================");
}

/* -------------------------- LVGL UI Initialization -------------------------- */
void lvgl_ui_init(lv_obj_t *parent) {
  // Set screen background to black
  lv_obj_set_style_bg_color(parent, lv_color_black(), 0);
  lv_obj_set_style_bg_opa(parent, LV_OPA_COVER, 0);

  // Title style (16px white, center)
  static lv_style_t title_style;
  lv_style_init(&title_style);
  lv_style_set_text_color(&title_style, lv_color_white());
  lv_style_set_text_font(&title_style, &lv_font_montserrat_16);
  lv_style_set_text_align(&title_style, LV_TEXT_ALIGN_CENTER);

  // Data style (14px white, left)
  static lv_style_t data_style;
  lv_style_init(&data_style);
  lv_style_set_text_color(&data_style, lv_color_white());
  lv_style_set_text_font(&data_style, &lv_font_montserrat_14);
  lv_style_set_text_align(&data_style, LV_TEXT_ALIGN_LEFT);

  // Status style (12px, center)
  static lv_style_t status_style;
  lv_style_init(&status_style);
  lv_style_set_text_font(&status_style, &lv_font_montserrat_12);
  lv_style_set_text_align(&status_style, LV_TEXT_ALIGN_CENTER);

  // Title label
  label_title = lv_label_create(parent);
  lv_obj_add_style(label_title, &title_style, 0);
  lv_label_set_text(label_title, "AXP2101 Power Monitor");
  lv_obj_align(label_title, LV_ALIGN_TOP_MID, 0, 15);

  // Battery Voltage (top-left)
  label_batt_vol = lv_label_create(parent);
  lv_obj_add_style(label_batt_vol, &data_style, 0);
  lv_label_set_text(label_batt_vol, "Batt Vol: -- mV");
  lv_obj_align(label_batt_vol, LV_ALIGN_TOP_LEFT, 15, 50);

  // VBUS Voltage (top-left)
  label_vbus_vol = lv_label_create(parent);
  lv_obj_add_style(label_vbus_vol, &data_style, 0);
  lv_label_set_text(label_vbus_vol, "VBUS Vol: -- mV");
  lv_obj_align(label_vbus_vol, LV_ALIGN_TOP_LEFT, 15, 85);

  // System Voltage (top-left)
  label_sys_vol = lv_label_create(parent);
  lv_obj_add_style(label_sys_vol, &data_style, 0);
  lv_label_set_text(label_sys_vol, "Sys Vol:  -- mV");
  lv_obj_align(label_sys_vol, LV_ALIGN_TOP_LEFT, 15, 120);

  // Battery Percentage (top-left)
  label_batt_per = lv_label_create(parent);
  lv_obj_add_style(label_batt_per, &data_style, 0);
  lv_label_set_text(label_batt_per, "Batt%%:    -- %%");
  lv_obj_align(label_batt_per, LV_ALIGN_TOP_LEFT, 15, 155);

  // Charging Status (top-left)
  label_charge = lv_label_create(parent);
  lv_obj_add_style(label_charge, &data_style, 0);
  lv_label_set_text(label_charge, "Charging: --");
  lv_obj_align(label_charge, LV_ALIGN_TOP_LEFT, 15, 190);

  // Global Status (bottom-center)
  label_status = lv_label_create(parent);
  lv_obj_add_style(label_status, &status_style, 0);
  lv_label_set_text(label_status, "Status: Initializing");
  lv_obj_align(label_status, LV_ALIGN_BOTTOM_MID, 0, -15);

  // Create timer (1000ms interval = 1s update)
  axp_timer = lv_timer_create(axp_data_callback, 1000, NULL);
}

/* -------------------------- Main Setup & Loop -------------------------- */
void setup() {
  // Initialize serial
  Serial.begin(115200);
  Serial.println("=== AXP2101 Power Monitor ===");

  // Initialize I2C
  Wire.begin(I2C_SDA_PIN, I2C_SCL_PIN);
  
  // Initialize MCP23017 (for LCD reset)
  if (!mcp.begin_I2C(0x20)) {
    Serial.println("ERROR: MCP23017 initialization failed!");
  } else {
    Serial.println("MCP23017 initialized at address 0x20");
    // Configure LCD reset pin as output
    mcp.pinMode(1, OUTPUT);
  }
  
  lcd_reset();
  
  // Initialize display
  if (!gfx->begin()) {
    Serial.println("Display initialization failed!");
    while (1);
  }
  gfx->fillScreen(RGB565_BLACK);

  // Turn on backlight
#ifdef GFX_BL
  pinMode(GFX_BL, OUTPUT);
  digitalWrite(GFX_BL, HIGH);
#endif

  // Initialize LVGL
  lv_init();
  screenWidth = LCD_HOR_RES;
  screenHeight = LCD_VER_RES;
  bufSize = screenWidth * 20; // Optimize buffer size

  // Allocate LVGL display buffers
  disp_draw_buf1 = (lv_color_t *)malloc(bufSize * sizeof(lv_color_t));
  disp_draw_buf2 = (lv_color_t *)malloc(bufSize * sizeof(lv_color_t));
  if (!disp_draw_buf1 || !disp_draw_buf2) {
    Serial.println("LVGL buffer allocation failed!");
    while (1);
  }
  lv_disp_draw_buf_init(&draw_buf, disp_draw_buf1, disp_draw_buf2, bufSize);

  // Register LVGL display driver
  lv_disp_drv_init(&disp_drv);
  disp_drv.hor_res = screenWidth;
  disp_drv.ver_res = screenHeight;
  disp_drv.flush_cb = my_disp_flush;
  disp_drv.draw_buf = &draw_buf;
  lv_disp_drv_register(&disp_drv);

  // Initialize AXP2101
  if (!init_sensors()) {
    Serial.println("Sensor initialization failed!");
    while (1);
  }

  // Initialize LVGL UI
  lvgl_ui_init(lv_scr_act());

  Serial.println("Setup complete - displaying AXP2101 data");
}

void loop() {
  lv_timer_handler(); // Process LVGL tasks
  delay(5);

  // Clear AXP2101 interrupts (prevent IRQ buffer overflow)
  power.clearIrqStatus();
}