#![no_std]
#![no_main]

use core::cell::RefCell;
use core::fmt::Write;
use embedded_hal::delay::DelayNs;

use esp_backtrace as _;
esp_bootloader_esp_idf::esp_app_desc!();
use esp_hal::{
    delay::Delay,
    gpio::{Level, Output, OutputConfig},
    i2c::master::{Config as I2cConfig, I2c},
    main,
    spi::master::{Config as SpiConfig, Spi},
    time::Rate,
};
use esp_println::println;

use embedded_hal_bus::i2c::RefCellDevice;

mod jd9853;
use jd9853::JD9853;
use mipidsi::{interface::SpiInterface, Builder};

use embedded_graphics::{
    mono_font::{
        ascii::{FONT_6X10, FONT_10X20},
        MonoTextStyle,
    },
    pixelcolor::Rgb565,
    prelude::*,
    primitives::{Line, PrimitiveStyle, PrimitiveStyleBuilder, Rectangle},
    text::Text,
};

use ina3221::INA3221;
use axs5106l::{Axs5106l, Rotation};
use heapless::String as HString;

// --- Hardware constants ------------------------------------------------------

const LCD_W: u16 = 172;
const LCD_H: u16 = 320;

//   Display pin  │ ESP32-C6 GPIO
//   ─────────────┼───────────────
//   SCL / SCK    │ GPIO1
//   SDA / MOSI   │ GPIO2
//   CS           │ GPIO14
//   DC           │ GPIO15
//   RST          │ GPIO22
//   BLK / BL     │ GPIO23 (set high to enable backlight)
//
//   I2C SDA      │ GPIO18
//   I2C SCL      │ GPIO19
//   Touch RST    │ GPIO20
//
//   INA3221 I2C address: 0x40 (A0 = GND)
//   Shunt resistor: 50 mΩ

/// Shunt resistor in milliohms.
const SHUNT_MOHMS: i32 = 50;

// --- App types ---------------------------------------------------------------

#[derive(Clone, Copy, PartialEq)]
enum AppState {
    Idle,
    Measuring,
}

#[derive(Clone, Copy, Default)]
struct Measurement {
    bus_uv: i32,     // bus voltage, µV
    current_ua: i32, // channel current, µA  (derived from shunt voltage)
    power_uw: i32,   // power, µW
}

// --- Formatting helpers ------------------------------------------------------

fn fmt_voltage(buf: &mut HString<24>, uv: i32) {
    let sign = if uv < 0 { "-" } else { "" };
    let uv = uv.unsigned_abs();
    write!(buf, "{}{}.{:02} V", sign, uv / 1_000_000, (uv % 1_000_000) / 10_000).ok();
}

fn fmt_current(buf: &mut HString<24>, ua: i32) {
    let sign = if ua < 0 { "-" } else { "" };
    let ua = ua.unsigned_abs();
    write!(buf, "{}{}.{} mA", sign, ua / 1000, (ua % 1000) / 100).ok();
}

fn fmt_power(buf: &mut HString<24>, uw: i32) {
    let sign = if uw < 0 { "-" } else { "" };
    let uw = uw.unsigned_abs();
    write!(buf, "{}{}.{} mW", sign, uw / 1000, (uw % 1000) / 100).ok();
}

fn fmt_duration(buf: &mut HString<24>, secs: u32) {
    let h = secs / 3600;
    let m = (secs % 3600) / 60;
    let s = secs % 60;
    write!(buf, "{:02}:{:02}:{:02}", h, m, s).ok();
}

// --- UI ----------------------------------------------------------------------

fn draw_ui<D: DrawTarget<Color = Rgb565>>(display: &mut D, state: AppState, m: &Measurement, elapsed_s: u32) {
    const BLUE:  Rgb565 = Rgb565::new(0, 0, 20);
    const GREEN: Rgb565 = Rgb565::new(0, 40, 0);
    const GREY:  Rgb565 = Rgb565::new(10, 20, 10);

    display.clear(Rgb565::BLACK).ok();

    // Header bar
    Rectangle::new(Point::new(0, 0), Size::new(LCD_W as u32, 42))
        .into_styled(PrimitiveStyleBuilder::new().fill_color(BLUE).build())
        .draw(display)
        .ok();
    Text::new(
        "POWER MONITOR",
        Point::new(21, 29),
        MonoTextStyle::new(&FONT_10X20, Rgb565::WHITE),
    )
    .draw(display)
    .ok();

    // Status line
    let (status_str, status_color) = if state == AppState::Measuring {
        ("* MEASURING", GREEN)
    } else {
        ("  STOPPED  ", GREY)
    };
    Text::new(status_str, Point::new(8, 68), MonoTextStyle::new(&FONT_10X20, status_color))
        .draw(display)
        .ok();

    // Divider
    Line::new(Point::new(0, 78), Point::new(LCD_W as i32 - 1, 78))
        .into_styled(PrimitiveStyle::with_stroke(GREY, 1))
        .draw(display)
        .ok();

    // Labels
    let label = MonoTextStyle::new(&FONT_6X10, GREY);
    Text::new("Voltage", Point::new(8, 100), label).draw(display).ok();
    Text::new("Current", Point::new(8, 160), label).draw(display).ok();
    Text::new("Power",   Point::new(8, 220), label).draw(display).ok();

    // Values
    let value = MonoTextStyle::new(&FONT_10X20, Rgb565::WHITE);
    let mut buf: HString<24> = HString::new();

    fmt_voltage(&mut buf, m.bus_uv);
    Text::new(&buf, Point::new(8, 128), value).draw(display).ok();
    buf.clear();

    fmt_current(&mut buf, m.current_ua);
    Text::new(&buf, Point::new(8, 188), value).draw(display).ok();
    buf.clear();

    fmt_power(&mut buf, m.power_uw);
    Text::new(&buf, Point::new(8, 248), value).draw(display).ok();

    // Footer
    Line::new(Point::new(0, 263), Point::new(LCD_W as i32 - 1, 263))
        .into_styled(PrimitiveStyle::with_stroke(GREY, 1))
        .draw(display)
        .ok();

    let mut buf: HString<24> = HString::new();
    write!(buf, "Runtime: ").ok();
    fmt_duration(&mut buf, elapsed_s);
    Text::new(&buf, Point::new(5, 278), MonoTextStyle::new(&FONT_6X10, GREY))
        .draw(display)
        .ok();

    Text::new(
        "tap to start / stop",
        Point::new(5, 293),
        MonoTextStyle::new(&FONT_6X10, GREY),
    )
    .draw(display)
    .ok();
}

// --- Entry point -------------------------------------------------------------

#[main]
fn main() -> ! {
    let peripherals = esp_hal::init(esp_hal::Config::default());

    esp_alloc::heap_allocator!(size: 72 * 1024);

    println!("Power monitor starting...");

    let mut delay = Delay::new();

    // --- SPI / display -------------------------------------------------------
    let spi = Spi::new(
        peripherals.SPI2,
        SpiConfig::default()
            .with_frequency(Rate::from_mhz(40))
            .with_mode(esp_hal::spi::Mode::_0),
    )
    .unwrap()
    .with_sck(peripherals.GPIO1)
    .with_mosi(peripherals.GPIO2);

    let cs  = Output::new(peripherals.GPIO14, Level::High, OutputConfig::default());
    let spi_device = embedded_hal_bus::spi::ExclusiveDevice::new(spi, cs, Delay::new()).unwrap();

    let dc  = Output::new(peripherals.GPIO15, Level::Low,  OutputConfig::default());
    let rst = Output::new(peripherals.GPIO22, Level::High, OutputConfig::default());
    let _bl = Output::new(peripherals.GPIO23, Level::High, OutputConfig::default());

    let mut spi_buf = [0u8; 512];
    let di = SpiInterface::new(spi_device, dc, &mut spi_buf);

    let mut display = Builder::new(JD9853, di)
        .reset_pin(rst)
        .display_size(LCD_W, LCD_H)
        .display_offset(34, 0)
        .init(&mut delay)
        .expect("display init failed");

    // --- I2C (shared between INA3221 and AXS5106L) ---------------------------
    let i2c = I2c::new(peripherals.I2C0, I2cConfig::default())
        .expect("I2C init failed")
        .with_sda(peripherals.GPIO18)
        .with_scl(peripherals.GPIO19);
    let i2c_bus = RefCell::new(i2c);

    // --- INA3221 (channel index 0 = physical CH1, shunt = 50 mΩ) ------------
    let ina = INA3221::new(RefCellDevice::new(&i2c_bus), 0x40);

    // --- AXS5106L touch controller -------------------------------------------
    let touch_rst = Output::new(peripherals.GPIO20, Level::High, OutputConfig::default());
    let mut touch = Axs5106l::new(
        RefCellDevice::new(&i2c_bus),
        touch_rst,
        LCD_W,
        LCD_H,
        Rotation::Rotate0,
    );
    touch.init(&mut delay).expect("touch init failed");

    // --- App state -----------------------------------------------------------
    let mut state         = AppState::Idle;
    let mut measurement   = Measurement::default();
    let mut was_touched   = false;
    let mut tick: u32     = 0;
    let mut runtime_loops: u32 = 0; // increments every 10 ms while measuring

    draw_ui(&mut display, state, &measurement, 0);

    loop {
        // Touch: toggle measuring on the leading edge of a tap
        if let Ok(td) = touch.read_raw_touch_data() {
            let touched = td.has_touches();
            if touched && !was_touched {
                if state == AppState::Idle {
                    state         = AppState::Measuring;
                    tick          = 50; // trigger first sensor read immediately
                    runtime_loops = 0;
                } else {
                    state = AppState::Idle;
                }
                draw_ui(&mut display, state, &measurement, runtime_loops / 100);
            }
            was_touched = touched;
        }

        // Read INA3221 every 500 ms while measuring (10 ms loop × 50 = 500 ms)
        if state == AppState::Measuring {
            runtime_loops = runtime_loops.saturating_add(1);
            tick += 1;
            if tick >= 50 {
                tick = 0;
                match (ina.get_shunt_voltage(0), ina.get_bus_voltage(0)) {
                    (Ok(shunt_v), Ok(bus_v)) => {
                        let bus_uv    = bus_v.micro_volts();
                        let shunt_uv  = shunt_v.micro_volts();
                        // I (µA) = V_shunt (µV) × 1000 / R_shunt (mΩ)
                        let current_ua = shunt_uv * 1000 / SHUNT_MOHMS;
                        // P (µW) = V_bus (µV) × I (µA) / 1 000 000
                        let power_uw =
                            ((bus_uv as i64 * current_ua as i64) / 1_000_000) as i32;
                        measurement = Measurement { bus_uv, current_ua, power_uw };
                        draw_ui(&mut display, state, &measurement, runtime_loops / 100);
                    }
                    _ => println!("INA3221 read error"),
                }
            }
        }

        delay.delay_ms(10u32);
    }
}
