#![no_std]
#![no_main]
#![deny(
    clippy::mem_forget,
    reason = "mem::forget is generally not safe to do with esp_hal types, especially those \
    holding buffers for the duration of a data transfer."
)]


use core::{num::{self, Wrapping}, ptr::addr_of_mut};

use alloc::{string::String, sync::Arc};
use embassy_executor::Spawner;
use embassy_time::{Instant, Timer};

use esp_hal::{gpio::{Output, OutputConfig, Pin}, time::Rate, xtensa_lx::debug_break};
use esp_hal::{
    clock::CpuClock, system::{AppCoreGuard, CpuControl, Stack}, timer::{systimer::SystemTimer, timg::{TimerGroup, Wdt}}
};

use embassy_sync::{
    pubsub::PubSubChannel,
    blocking_mutex::raw::NoopRawMutex
};
use static_cell::ConstStaticCell;
use log::*;
use renderbug_embassy::{events::Prediction, graphics::display::DisplayControls, logging::RenderbugLogger, tasks::{oled::{OledUI, OledUiSurfacePool, oled_ui}, safetyui::{SafetyUi, safety_ui_main}, ui::UiSurfacePool}};
use renderbug_embassy::events::Measurement;
use static_cell::StaticCell;
use esp_backtrace as _;
use esp_rtos::embassy::{Executor, InterruptExecutor};
use esp_hal::interrupt::software::SoftwareInterruptControl;

use renderbug_embassy::tasks::{
    motion::motion_task,
    ui::{Ui, ui_main}
};

use embassy_sync::{blocking_mutex::raw::CriticalSectionRawMutex, pubsub::DynSubscriber};
use embassy_sync::channel::Channel;

extern crate alloc;

// This creates a default app-descriptor required by the esp-idf bootloader.
// For more information see: <https://docs.espressif.com/projects/esp-idf/en/stable/esp32/api-reference/system/app_image_format.html#application-description>
esp_bootloader_esp_idf::esp_app_desc!();

#[cfg(feature="radio")]
static WIFI_INIT: StaticCell<esp_radio::Controller<'static>> = StaticCell::new();

#[esp_rtos::main]
async fn main(spawner: Spawner) {
    esp_alloc::heap_allocator!(size: 128 * 1024);
    RenderbugLogger::init_logger();

    let config = esp_hal::Config::default().with_cpu_clock(CpuClock::max());
    let peripherals = esp_hal::init(config);

    let sys_timer = SystemTimer::new(peripherals.SYSTIMER);
    esp_rtos::start(sys_timer.alarm0);
    info!("Embassy initialized!");

    static MOTION_BUS: StaticCell<Channel<CriticalSectionRawMutex,Measurement,5> > = StaticCell::new();
    let motion_bus = MOTION_BUS.init_with(|| { Channel::new() });

    info!("Setting up rendering pipeline");
    let mut surfaces = UiSurfacePool::default();
    let ui = Ui::new(&mut surfaces);

    let display_controls = DisplayControls::default();
    let oled_controls = DisplayControls::default();

    let mut oled_surfaces = OledUiSurfacePool::default();
    let oled_uniforms = Default::default();
    let oledui = OledUI::new(&mut oled_surfaces, oled_controls.clone(), Arc::clone(&oled_uniforms));

    let mut safety_surfaces = UiSurfacePool::default();
    let safety_ui = SafetyUi::new(&mut safety_surfaces, display_controls.clone());

    let timer0 = TimerGroup::new(peripherals.TIMG0);
    let mut wdt = timer0.wdt;
    wdt.set_timeout(esp_hal::timer::timg::MwdtStage::Stage0, esp_hal::time::Duration::from_secs(5));
    wdt.enable();

    // Spawn the rendering task as soon as possible so it can start pushing pixels
    spawner.must_spawn(renderbug_bike::tasks::render::render(peripherals.SPI2.degrade(), peripherals.DMA_CH2.degrade(), peripherals.GPIO5.degrade(), surfaces, safety_surfaces, display_controls, wdt));

    #[cfg(feature="motion")]
    {
        use embassy_embedded_hal::shared_bus::asynch::i2c::I2cDevice;
        use esp_hal::{i2c::master::{Config, I2c}, Async};
        use embassy_sync::{blocking_mutex::raw::NoopRawMutex, mutex::Mutex};

        static I2C_BUS: StaticCell<Mutex<NoopRawMutex, I2c<'static, Async>>> = StaticCell::new();

        info!("Launching i2c sensor tasks");
        let sda = peripherals.GPIO4;
        let scl = peripherals.GPIO3;
        let i2c = I2c::new(peripherals.I2C1, Config::default()).unwrap().with_scl(scl).with_sda(sda).into_async();
        let i2c_bus = I2C_BUS.init(Mutex::new(i2c));
        #[cfg(feature="mpu")]
        spawner.must_spawn(renderbug_embassy::tasks::mpu::mpu_task(motion_bus.dyn_sender(), I2cDevice::new(i2c_bus)));
        #[cfg(feature="gps")]
        spawner.must_spawn(renderbug_embassy::tasks::gps::gps_task(motion_bus.dyn_sender(), I2cDevice::new(i2c_bus)));
    }

    #[cfg(feature="oled")]
    {
        use esp_hal::i2c::master::{Config, I2c};
        use renderbug_embassy::graphics::ssd1306::SsdOutput;

        let rst = Output::new(peripherals.GPIO21, esp_hal::gpio::Level::Low, OutputConfig::default());
        let i2c = I2c::new(
            peripherals.I2C0, 
            Config::default().with_frequency(Rate::from_khz(400))
        ).unwrap().with_scl(peripherals.GPIO18).with_sda(peripherals.GPIO17).into_async();
        let output = SsdOutput::new(i2c, rst, oled_controls).await;
        spawner.must_spawn(renderbug_embassy::tasks::oled_render::oled_render(output, oled_surfaces, oled_uniforms));
    }

    #[cfg(feature="simulation")]
    {
        use esp_storage::FlashStorage;
        use renderbug_embassy::tasks::simulation::{SharedFlash, SimDataTable};
        let mut storage = SharedFlash::new(FlashStorage::new());
        let mut buf = [8; 1024];
        let partitions = esp_bootloader_esp_idf::partitions::read_partition_table(&mut storage, &mut buf).unwrap();
        for sim_data in SimDataTable::open(storage, partitions).expect("Could not find sim data!") {
            let srcid = sim_data.srcid();
            info!("Found simulation data for {srcid:?}");
            if spawner.spawn(renderbug_embassy::tasks::simulation::simulation_task(sim_data, motion_bus.dyn_sender())).is_err() {
                error!("Unable to spawn simulation task for {srcid:?}! Increase the task pool size.");
            }
        }
    }

    #[cfg(feature="radio")]
    let (wifi, network_device, ble) = {
        info!("Configuring wifi");
        esp_radio::wifi_set_log_verbose();
        let wifi_init = WIFI_INIT.init_with(|| {esp_radio::init().expect("Failed to initialize radio controller")});

        let ble = esp_radio::ble::controller::BleConnector::new(wifi_init, peripherals.BT, esp_radio::ble::Config::default()).unwrap();

        let (wifi, interfaces) = esp_radio::wifi::new(wifi_init, peripherals.WIFI, esp_radio::wifi::Config::default())
            .expect("Failed to initialize WIFI!");

        (wifi, interfaces.sta, ble)
    };

    let core2_main = |spawner: Spawner| {
        info!("Starting application tasks");

        static PREDICTIONS: StaticCell<PubSubChannel<NoopRawMutex, Prediction, 15, 6, 1>> = StaticCell::new();
        let predictions = PREDICTIONS.init(PubSubChannel::new());

        #[cfg(not(feature="demo"))]
        {
            info!("Launching motion engine");
            spawner.must_spawn(motion_task(motion_bus.dyn_receiver(), predictions.dyn_publisher().unwrap()));
        }

        #[cfg(feature="demo")]
        {
            warn!("Launching with demo sequencer");
            spawner.must_spawn(renderbug_embassy::tasks::demo::demo_task(predictions.dyn_publisher().unwrap()));
        }
        
        info!("Launching Safety UI");
        spawner.must_spawn(safety_ui_main(predictions.dyn_subscriber().unwrap(), safety_ui));
        info!("Launching UI");
        spawner.must_spawn(ui_main(predictions.dyn_subscriber().unwrap(), ui));
        info!("Launching OLED UI");
        spawner.must_spawn(oled_ui(predictions.dyn_subscriber().unwrap(), oledui));

        #[cfg(feature="radio")]
        {
            use embassy_net::StackResources;
            use esp_hal::rng::Rng;
            use static_cell::ConstStaticCell;

            info!("Setting up networking stack");
            static RESOURCES: ConstStaticCell<StackResources<5>> = ConstStaticCell::new(StackResources::new());
            let config = embassy_net::Config::dhcpv4(Default::default());
            let seed = Rng::new().random() as i32;
            let (stack, runner) = embassy_net::new(network_device, config, RESOURCES.take(), seed as u64);
            info!("Launching network services");
            //spawner.must_spawn(renderbug_embassy::tasks::wifi::net_task(runner));

            info!("Starting connectivity task");
            //spawner.must_spawn(renderbug_embassy::tasks::wifi::wifi_connect_task(wifi, stack, motion_bus.dyn_sender()));

            info!("Launching HTTP telemetry");
            //spawner.must_spawn(renderbug_embassy::tasks::wifi::http_telemetry_task(predictions.dyn_subscriber().unwrap(), stack));

            info!("Starting BLE services");
            spawner.must_spawn(renderbug_embassy::tasks::ble::ble_task(ble, predictions.dyn_subscriber().unwrap(), spawner));
        }

        #[cfg(feature="dual-core")]
        {
            info!("Launching core 2 watchdog");
            let timer1 = TimerGroup::new(peripherals.TIMG1);
            let mut ui_wdt = timer1.wdt;
            ui_wdt.set_timeout(esp_hal::timer::timg::MwdtStage::Stage0, esp_hal::time::Duration::from_secs(10));
            #[cfg(feature="dual-core")]
            ui_wdt.enable();
            spawner.must_spawn(wdt_task(ui_wdt));
        }

        spawner.must_spawn(print_telemetry(predictions.dyn_subscriber().unwrap()));

        info!("System is ready in {}ms", Instant::now().as_millis());
    };

    #[cfg(feature="dual-core")]
    {
        static mut CORE2_STACK: Stack<16384> = Stack::new();
        let swi = SoftwareInterruptControl::new(peripherals.SW_INTERRUPT);
        esp_rtos::start_second_core(peripherals.CPU_CTRL, swi.software_interrupt0, swi.software_interrupt1, unsafe { &mut *addr_of_mut!(CORE2_STACK) }, || {
            info!("Second CPU core started");
            static CORE2_EXEC: StaticCell<Executor> = StaticCell::new();
            let exec = CORE2_EXEC.init_with(|| { Executor::new() });
            exec.run(core2_main);
        });
    }

    #[cfg(not(feature="dual-core"))]
    core2_main(spawner);

    info!("Ready to rock and roll");
}

#[embassy_executor::task]
async fn wdt_task(mut wdt: Wdt<esp_hal::peripherals::TIMG1<'static>>) {
    loop {
        // Watchdog is set to trip after 5 seconds, so we wait just long enough before it trips
        // TODO: We should maybe extend this timeout to 30s or 1m when the system is sleeping
        trace!("Bark wdt");
        wdt.feed();
        Timer::after_secs(3).await;
    }
}

#[embassy_executor::task]
async fn print_telemetry(mut events: DynSubscriber<'static, Prediction>) {
    info!("telemetry ready");
    let mut num_events = Wrapping(0usize);
    loop {
        let next = events.next_message_pure().await;
        trace!("idx={} predict={next:?}", num_events.0);
        num_events += 1;        
    }
}