//! embassy hello world
//!
//! This is an example of running the embassy executor with multiple tasks
//! concurrently.

#![no_std]
#![no_main]
#![feature(type_alias_impl_trait)]
#![feature(trait_alias)]
#![feature(generic_arg_infer)]
extern crate alloc;
use core::mem::{self, MaybeUninit};

use alloc::borrow::Cow;
use alloc::boxed::Box;
use alloc::collections::BTreeMap;
use alloc::string::{String, ToString};
use embassy_executor::Executor;
use embassy_net::{Config, DhcpConfig, Stack, StackResources};
use embassy_sync::blocking_mutex::raw::{CriticalSectionRawMutex, NoopRawMutex};
use embassy_sync::mutex::Mutex;
use embassy_time::{Duration, Timer};
use esp_alloc::EspHeap;
use esp_backtrace as _;
use esp_hal::adc::{AdcPin, Attenuation, ADC};
use esp_hal::clock::CpuClock;
use esp_hal::gpio::{
    Analog, AnalogPin, AnyPin, GpioPin, OpenDrain, Output, OutputSignal, PushPull,
};
use esp_hal::peripherals::{Peripherals, ADC1, ADC2, TIMG0};
use esp_hal::systimer::SystemTimer;
use esp_hal::{adc::AdcConfig, clock::ClockControl};
use esp_hal::{embassy, prelude::*, rmt, Rmt, Rng, Rtc, IO};
use esp_hal_smartled::SmartLedsAdapter;
use esp_hal_smartled::*;
use esp_println::println;
use esp_wifi::wifi::{get_random, ClientConfiguration, Configuration};
use esp_wifi::wifi::{WifiController, WifiDevice, WifiEvent, WifiStaDevice, WifiState};
use esp_wifi::{initialize as initialize_wifi, EspWifiInitFor};
use smart_leds::{SmartLedsWrite, RGB8};
use static_cell::{make_static, StaticCell};

use crate::mqtt::publish_data;

mod mqtt;

#[global_allocator]
static ALLOCATOR: EspHeap = EspHeap::empty();

const SSID: Option<&str> = option_env!("WIFI_SSID");
const PASSWORD: Option<&str> = option_env!("WIFI_PASSWORD");

pub type NetworkStack = &'static Stack<WifiDevice<'static, WifiStaDevice>>;

fn init_heap() {
    const HEAP_SIZE: usize = 32 * 1024;
    static mut HEAP: MaybeUninit<[u8; HEAP_SIZE]> = MaybeUninit::uninit();

    unsafe {
        ALLOCATOR.init(HEAP.as_mut_ptr() as *mut u8, HEAP_SIZE);
    }
}
trait MySmartLed {
    fn set_color(&mut self, color: RGB8);
}

impl<TX: rmt::TxChannel, const BUFFER_SIZE: usize> MySmartLed
    for SmartLedsAdapter<TX, BUFFER_SIZE>
{
    fn set_color(&mut self, color: RGB8) {
        if let Err(err) = self.write([color].into_iter()) {
            esp_println::println!("{err:?}");
        }
    }
}

pub static DATA: Mutex<CriticalSectionRawMutex, BTreeMap<Cow<'static, str>, String>> =
    Mutex::new(BTreeMap::new());

#[embassy_executor::task]
async fn blink(mut led: Box<dyn MySmartLed>) {
    loop {
        esp_println::println!("Bing!");
        let scale = 4;
        for r in 0..(255 / scale) {
            Timer::after(Duration::from_millis(1)).await;
            for g in 0..(255 / scale) {
                Timer::after(Duration::from_millis(1)).await;
                for b in 0..(255 / scale) {
                    Timer::after(Duration::from_millis(1)).await;
                    let color = RGB8::new(
                        (r * scale % 255) as _,
                        (g * scale % 255) as _,
                        (b * scale % 255) as _,
                    );
                    led.set_color(color);
                }
            }
        }
    }
}

async fn async_read<T, E>(
    mut fun: impl FnMut() -> nb::Result<T, E>,
    interval: Duration,
) -> Result<T, E> {
    loop {
        match fun() {
            Err(nb::Error::WouldBlock) => Timer::after(interval).await,
            Err(nb::Error::Other(e)) => return Err(e),
            Ok(val) => return Ok(val),
        }
    }
}

/// https://docs.espressif.com/projects/esp-idf/en/v4.4/esp32c3/api-reference/peripherals/adc.html
fn adc_voltage(d_out: u16, att: Attenuation) -> u16 {
    let scale = match att {
        Attenuation::Attenuation0dB => 800,
        Attenuation::Attenuation6dB => 1350,
        Attenuation::Attenuation2p5dB => 1100,
        Attenuation::Attenuation11dB => 2600,
    };
    (d_out as u32 * scale / 4095) as u16
}

#[embassy_executor::task]
async fn moisture(
    mut pin: AdcPin<GpioPin<Analog, 4>, ADC1>,
    adc: &'static Mutex<NoopRawMutex, ADC<'static, ADC1>>,
    mut supply: AnyPin<Output<PushPull>>,
) {
    let warmup = Duration::from_millis(100);
    let submerged_in_water = 1500;
    let dry = 2600;
    loop {
        supply.set_high().unwrap();
        Timer::after(warmup).await;
        let pin_value: u16 = {
            let mut adc1 = adc.lock().await;
            async_read(|| adc1.read(&mut pin), Duration::from_millis(10)).await
        }
        .unwrap();

        supply.set_low();
        // Vout = Dout * Vmax / Dmax
        let milli_volt = adc_voltage(pin_value, Attenuation::Attenuation11dB) as u32;
        let moisture = ((milli_volt.checked_sub(submerged_in_water).unwrap_or(0))
            << 8 / (dry - submerged_in_water))
            >> 8;
        esp_println::println!("moisture: {moisture}%, v_s: {milli_volt}");
        {
            DATA.lock()
                .await
                .insert(Cow::from("moisture"), moisture.to_string());
        }
        Timer::after(Duration::from_secs(10) - warmup).await;
    }
}

#[embassy_executor::task]
async fn battery_monitor(
    mut vbat_in: AdcPin<GpioPin<Analog, 3>, ADC1>,
    adc: &'static Mutex<NoopRawMutex, ADC<'static, ADC1>>,
) {
    loop {
        let v_out = adc_voltage(
            {
                let mut adc1 = adc.lock().await;
                async_read(|| adc1.read(&mut vbat_in), Duration::from_millis(10)).await
            }
            .unwrap(),
            Attenuation::Attenuation11dB,
        );
        // account for 50:50 voltage divider
        let v_bat = v_out * 2;
        println!("V_bat: {}", v_bat);
        {
            DATA.lock()
                .await
                .insert(Cow::from("vbat"), v_bat.to_string());
        }
        Timer::after(Duration::from_secs(15)).await;
    }
}
static EXECUTOR: StaticCell<Executor> = StaticCell::new();

#[entry]
fn main() -> ! {
    esp_println::println!("Init!");
    init_heap();

    let peripherals = Peripherals::take();
    let system = peripherals.SYSTEM.split();
    let clocks = ClockControl::configure(system.clock_control, CpuClock::Clock160MHz).freeze();

    let mut rtc = Rtc::new(peripherals.LPWR);
    let mut timer_group0 = esp_hal::timer::TimerGroup::new(peripherals.TIMG0, &clocks);
    let timer_group1 = esp_hal::timer::TimerGroup::new(peripherals.TIMG1, &clocks);
    {
        let mut wdt0 = &mut timer_group0.wdt;
        let mut wdt1 = timer_group1.wdt;

        // Disable watchdog timers
        rtc.swd.disable();
        rtc.rwdt.disable();
        wdt0.disable();
        wdt1.disable();
    }
    let timer = SystemTimer::new(peripherals.SYSTIMER).alarm0;

    let mut rng = Rng::new(peripherals.RNG);

    let mut seed = [0u8; 32]; // very random, very secure seed
    rng.read(&mut seed).expect("random seed");

    let wifi_init = initialize_wifi(
        EspWifiInitFor::Wifi,
        timer,
        rng,
        system.radio_clock_control,
        &clocks,
    )
    .unwrap();
    let wifi = peripherals.WIFI;
    let (wifi_interface, controller) =
        esp_wifi::wifi::new_with_mode(&wifi_init, wifi, WifiStaDevice).unwrap();

    let mut hostname = heapless::String::<_>::new();
    hostname.push_str("esp32c3");
    let config = Config::dhcpv4({
        let mut dhcp = DhcpConfig::default();
        dhcp.hostname = Some(hostname);
        dhcp
    });
    // Init network stack
    let stack = &*make_static!(Stack::new(
        wifi_interface,
        config,
        make_static!(StackResources::<3>::new()),
        u64::from_le_bytes(seed[..8].try_into().unwrap())
    ));

    embassy::init(&clocks, timer_group0);

    let io = IO::new(peripherals.GPIO, peripherals.IO_MUX);
    let rmt = Rmt::new(peripherals.RMT, 80u32.MHz(), &clocks).unwrap();

    let rmt_buffer = smartLedBuffer!(1);
    let mut led = SmartLedsAdapter::new(rmt.channel0, io.pins.gpio7, rmt_buffer, &clocks);

    led.set_color(RGB8::new(0, 255, 255));

    let led: Box<dyn MySmartLed> = Box::new(led);

    let mut adc1_config = AdcConfig::new();
    let pin = adc1_config.enable_pin(io.pins.gpio4.into_analog(), Attenuation::Attenuation11dB);
    let vbat_in = adc1_config.enable_pin(io.pins.gpio3.into_analog(), Attenuation::Attenuation11dB);
    let moisture_sensor_suppy_pin = io.pins.gpio5.into_push_pull_output().degrade();

    let adc1: &'static Mutex<NoopRawMutex, ADC<'static, ADC1>> =
        make_static!(Mutex::new(ADC::<ADC1>::new(peripherals.ADC1, adc1_config)));

    let executor = EXECUTOR.init(Executor::new());
    executor.run(move |spawner| {
        if let (Some(ssid), Some(psk)) = (SSID.as_ref(), PASSWORD.as_ref()) {
            spawner
                .spawn(wifi_connection(controller, ssid, psk))
                .unwrap();
            spawner.spawn(net_task(&stack)).unwrap();
            spawner.spawn(ip_task(&stack)).unwrap();
        }
        spawner.spawn(blink(led)).unwrap();
        spawner
            .spawn(moisture(pin, adc1, moisture_sensor_suppy_pin.into()))
            .unwrap();
        spawner.spawn(battery_monitor(vbat_in, adc1)).unwrap();
        spawner.spawn(publish_data(&stack, seed, Duration::from_secs(60)));
    })
}

#[embassy_executor::task]
async fn wifi_connection(
    mut controller: WifiController<'static>,
    ssid: &'static str,
    psk: &'static str,
) {
    println!("start connection task");
    println!("Device capabilities: {:?}", controller.get_capabilities());
    loop {
        match esp_wifi::wifi::get_wifi_state() {
            WifiState::StaConnected => {
                // wait until we're no longer connected
                controller.wait_for_event(WifiEvent::StaDisconnected).await;
                Timer::after(Duration::from_millis(5000)).await
            }
            _ => {}
        }
        if !matches!(controller.is_started(), Ok(true)) {
            let client_config = Configuration::Client(ClientConfiguration {
                ssid: ssid.try_into().unwrap(),
                password: psk.try_into().unwrap(),
                ..Default::default()
            });
            controller.set_configuration(&client_config).unwrap();
            println!("Starting wifi");
            controller.start().await.unwrap();
            println!("Wifi started!");
        }
        println!("About to connect...");

        match controller.connect().await {
            Ok(_) => println!("Wifi connected!"),
            Err(e) => {
                println!("Failed to connect to wifi: {e:?}");
                Timer::after(Duration::from_millis(5000)).await
            }
        }
    }
}

#[embassy_executor::task]
async fn ip_task(stack: NetworkStack) {
    loop {
        if stack.is_link_up() {
            break;
        }
        Timer::after(Duration::from_millis(500)).await;
    }
    println!("Waiting to get IP address...");
    loop {
        if let Some(config) = stack.config_v4() {
            println!("Got IP: {}", config.address);
            break;
        }
        Timer::after(Duration::from_millis(500)).await;
    }
}

#[embassy_executor::task]
async fn net_task(stack: NetworkStack) {
    stack.run().await
}