use anyhow::{Result, anyhow};
use esp_idf_svc::hal::gpio;
use esp_idf_svc::hal::delay::FreeRtos;
use esp_idf_svc::hal::i2c::{I2cConfig, I2cDriver};
use esp_idf_svc::hal::peripherals::Peripherals;
use esp_idf_svc::hal::delay::BLOCK;
use esp_idf_svc::hal::prelude::FromValueType;
use esp_idf_svc::hal::units::Hertz;
use esp_idf_svc::hal::uart::{config, UartDriver};

use embedded_graphics::{
    mono_font::{
        ascii::{FONT_10X20},
        MonoTextStyle,
    },
    image::{Image, ImageRaw},
    pixelcolor::BinaryColor,
    prelude::*,
    text::Text,
};
use ssd1306::{I2CDisplayInterface, Ssd1306, size::*, rotation::*};
use ssd1306::mode::DisplayConfig;

use std::thread;
use std::sync::{Arc, Mutex};
use log::{info, warn};

#[derive(Debug, Clone, Default)]
pub struct Pms7003Data {
    // PM concentrations (CF=1, standard particle) in μg/m³
    pub pm1_0_cf1: u16,
    pub pm2_5_cf1: u16,
    pub pm10_cf1: u16,

    // PM concentrations (atmospheric environment) in μg/m³
    pub pm1_0_atm: u16,
    pub pm2_5_atm: u16,
    pub pm10_atm: u16,

    // Particle counts (particles per 0.1L air)
    pub particles_0_3um: u16,
    pub particles_0_5um: u16,
    pub particles_1_0um: u16,
    pub particles_2_5um: u16,
    pub particles_5_0um: u16,
    pub particles_10um: u16,
}

// u16::from_be_bytes is too verbose
fn get_u16(buf: &[u8]) -> u16 { ((buf[0] as u16) << 8) | buf[1] as u16 }

fn read_pm_data(buf: &[u8; 32]) -> Option<Pms7003Data> {
    let checksum: u16 = buf.iter().take(29).copied().map(u16::from).sum();
    let csum_target = get_u16(&buf[30..32]);
    if checksum != csum_target {
        warn!("Checksum from PMS7003 READ does not match, skipping read");
        return None
    }
    Some(Pms7003Data {
        pm1_0_cf1:       get_u16(&buf[ 4..6 ]),
        pm2_5_cf1:       get_u16(&buf[ 6..8 ]),
        pm10_cf1:        get_u16(&buf[ 8..10]),
        pm1_0_atm:       get_u16(&buf[10..12]),
        pm2_5_atm:       get_u16(&buf[12..14]),
        pm10_atm:        get_u16(&buf[14..16]),
        particles_0_3um: get_u16(&buf[16..18]),
        particles_0_5um: get_u16(&buf[18..20]),
        particles_1_0um: get_u16(&buf[20..22]),
        particles_2_5um: get_u16(&buf[22..24]),
        particles_5_0um: get_u16(&buf[24..26]),
        particles_10um:  get_u16(&buf[26..28]),
    })
}

fn send_pms7003_command(uart: &UartDriver, cmd: u8, datah: u8, datal: u8) -> Result<()> {
    let mut command = [0x42u8, 0x4D, cmd, datah, datal, 0, 0];

    let mut checksum: u16 = 0;
    for i in 0..5 {
        checksum += command[i] as u16;
    }
    command[5] = ((checksum >> 8) & 0xFF) as u8;
    command[6] = (checksum & 0xFF) as u8;

    uart.write(&command)?;
    Ok(())
}

fn pm_sensor_task(pm_data: Arc<Mutex<Pms7003Data>>, uart: &UartDriver) -> Result<()> {
    info!("Staring UART PM Sensor Task");

    send_pms7003_command(uart, 0xE1, 0x00, 0x00)?;

    loop {
        send_pms7003_command(uart, 0xE2, 0x00, 0x00)?;

        let mut buf = [0_u8; 32];
        let len = uart.read(&mut buf, BLOCK)?;
        if len > 0 {
            if let Some(data) = read_pm_data(&buf) {
                let mut pm_data = pm_data.lock().unwrap();
                *pm_data = data;
            }
        } else {
            // info!("No bytes read");
        }
        uart.clear_rx()?;
        FreeRtos::delay_ms(5000);
    }
}

fn oled_task(pm_data: Arc<Mutex<Pms7003Data>>, i2c: I2cDriver) -> Result<()> {
    info!("Staring OLED Task");

    let interface = I2CDisplayInterface::new(i2c);
    let mut display = Ssd1306::new(interface, DisplaySize128x64, DisplayRotation::Rotate0)
        .into_buffered_graphics_mode();

    display.init().unwrap();

    let raw: ImageRaw<BinaryColor> = ImageRaw::new(include_bytes!("../rust.raw"), 64);
    let im = Image::new(&raw, Point::new(32, 0));

    im.draw(&mut display).unwrap();
    display.flush().map_err(|e| anyhow!("Could not flush display! {:?}", e))?;

    FreeRtos::delay_ms(2000);
    loop {
        let current_pm_data = pm_data.lock().unwrap().clone();
        display.clear(BinaryColor::Off).map_err(|e| anyhow!("Could not clear display! {:?}", e))?;
        Text::new(
            &format!("PM 2.5: {}", current_pm_data.pm2_5_atm),
            Point::new(20, 40),
            MonoTextStyle::new(&FONT_10X20, BinaryColor::On),
        ).draw(&mut display).map_err(|e| anyhow!("Could not draw text! {:?}", e))?;
        display.flush().map_err(|e| anyhow!("Could not flush display! {:?}", e))?;
        FreeRtos::delay_ms(500);
    }
}

fn main() -> Result<()> {
    // It is necessary to call this function once. Otherwise some patches to the runtime
    // implemented by esp-idf-sys might not link properly. See https://github.com/esp-rs/esp-idf-template/issues/71
    esp_idf_svc::sys::link_patches();

    // Initialize NVS
    // esp_idf_hal::nvs::EspDefaultNvsPartition::take()
    //     .context("Could not take NVS partition")?;
    esp_idf_svc::log::EspLogger::initialize_default();

    let peripherals = Peripherals::take()?;
    let i2c = peripherals.i2c0;
    let sda = peripherals.pins.gpio21;
    let scl = peripherals.pins.gpio22;

    let config = I2cConfig::new().baudrate(100.kHz().into());
    let i2c = I2cDriver::new(i2c, sda, scl, &config)?;

    let tx = peripherals.pins.gpio17;
    let rx = peripherals.pins.gpio16;

    let config = config::Config::new().baudrate(Hertz(9600));
    let uart = UartDriver::new(
        peripherals.uart2,
        tx,
        rx,
        Option::<gpio::Gpio0>::None,
        Option::<gpio::Gpio1>::None,
        &config,
    )?;

    let pm_data = Arc::new(Mutex::new(Pms7003Data::default()));

    let pm_data1 = Arc::clone(&pm_data);
    let pm_data2 = Arc::clone(&pm_data);

    let handles = vec! [
        thread::spawn(move || {
            if let Err(e) = oled_task(pm_data1, i2c) {
                log::error!("OLED Task Error: {:?}", e);
            }
        }),
        thread::spawn(move || {
            if let Err(e) = pm_sensor_task(pm_data2, &uart) {
                log::error!("PM Sensor Task Error: {:?}", e);
            }
        }),
    ];
    for h in handles {
        h.join().unwrap();
    }

    Ok(())
}