//! TODO:
//! * GUI event dispatch.
//! * Async interrupt handling of keyboard input. Reduces LCD glitching.
//! * Attempt to reduce the size of the framebuffer to 240x960 by changing the front/back porch of
//!   the LCD driver. Reduces LCD glitching.
//! * Bounce buffers to get rid of glitching completely.
//!   https://esp32.com/viewtopic.php?t=28230
//!   https://esp32.com/viewtopic.php?f=12&t=26793&start=20#p95677
#![no_std]
#![no_main]
#![feature(allocator_api)]
#![feature(btreemap_alloc)]
#![feature(macro_metavar_expr)]
#![feature(c_variadic)]
#![feature(c_size_t)]
#![feature(debug_closure_helpers)]

extern crate alloc;

use core::cell::RefCell;
use core::sync::atomic::{AtomicBool, Ordering};

use alloc::boxed::Box;
use alloc::collections::vec_deque::VecDeque;
use alloc::format;
use alloc::sync::Arc;
use embassy_executor::Spawner;
use embassy_sync::blocking_mutex;
use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
use embassy_sync::channel::Channel;
use embassy_time::{Duration, Timer};
use esp_alloc::MemoryCapability;
use esp_hal::clock::CpuClock;
use esp_hal::dma::{BurstConfig, ExternalBurstConfig, InternalBurstConfig};
use esp_hal::efuse::Efuse;
use esp_hal::gpio::{Input, InputConfig, Level, Output, OutputConfig, Pull};
use esp_hal::i2c::master::{I2c, I2cAddress};
use esp_hal::interrupt::software::{SoftwareInterrupt, SoftwareInterruptControl};
use esp_hal::psram::{FlashFreq, PsramConfig, PsramSize, SpiRamFreq, SpiTimingConfigCoreClock};
use esp_hal::rng::TrngSource;
use esp_hal::sha::ShaBackend;
use esp_hal::system::Stack;
use esp_hal::timer::timg::TimerGroup;
use esp_hal::uart::{Uart, UartRx};
use esp_hal::{Blocking, interrupt};
use esp_hal_bounce_buffers::{DmaBounce, RunningDmaBounceHandle};
use esp_rtos::embassy::{Executor, InterruptExecutor};
use esp_storage::FlashStorage;
use itertools::Itertools;
use log::{info, warn};
use rmk::channel::{CONTROLLER_CHANNEL, ControllerSub};
use rmk::config::{DeviceConfig, RmkConfig, StorageConfig, VialConfig};
use rmk::controller::{Controller, EventController};
use rmk::debounce::default_debouncer::DefaultDebouncer;
use rmk::event::ControllerEvent;
use rmk::hid::Report;
use rmk::input_device::Runnable;
use rmk::join_all;
use rmk::keyboard::Keyboard;
use rmk::types::action::{Action, KeyAction};
use rmk::{initialize_keymap_and_storage, run_devices, run_rmk};
use static_cell::StaticCell;
use {esp_alloc as _, esp_backtrace as _};

use crate::dpi::{DisplayPeripherals, Framebuffer};
use crate::flash::Partition;
use crate::matrix::IoeMatrix;
use crate::peripherals::st7701s::St7701sController;
use crate::proxy::create_hid_report_interceptor;
use crate::ram::{PSRAM_ALLOCATOR, STACK_SIZE_CORE_APP};
use crate::ui::backend::SlintBackend;
use crate::vial::{
    CustomKeycodes, VIAL_KEYBOARD_DEF, VIAL_KEYBOARD_ID, VIAL_KEYBOARD_NAME, VIAL_PRODUCT_ID,
    VIAL_VENDOR_ID,
};

mutually_exclusive_features::none_or_one_of!["usb-log", "alt-log", "rtt-log"];

mod config;
mod crypto;
mod db;
mod dpi;
mod ffi;
mod flash;
mod logging;
mod matrix;
mod peripherals;
mod proxy;
mod ram;
mod ui;
mod util;
mod vial;

#[cfg(feature = "alt-log")]
mod console;

// 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!();

// const FRAME_DURATION_MIN: Duration = Duration::from_millis(40); // 25 FPS
const FRAME_DURATION_MIN: Duration = Duration::from_millis(100); // 10 FPS

static KEYBOARD_REPORT_PROXY: Channel<CriticalSectionRawMutex, Report, 16> = Channel::new();
static LCD_ENABLED: AtomicBool = AtomicBool::new(false);

#[esp_rtos::main]
async fn main(_spawner: Spawner) {
    let config = esp_hal::Config::default()
        .with_cpu_clock(CpuClock::max())
        .with_psram(PsramConfig {
            size: PsramSize::AutoDetect,
            core_clock: Some(SpiTimingConfigCoreClock::SpiTimingConfigCoreClock80m),
            flash_frequency: FlashFreq::FlashFreq120m,
            ram_frequency: SpiRamFreq::Freq120m,
        });
    let peripherals: esp_hal::peripherals::Peripherals = esp_hal::init(config);

    #[allow(unused)]
    let (uart_rx, uart_tx) = {
        #[cfg(feature = "alt-log")]
        let (tx, rx) = (peripherals.GPIO12, peripherals.GPIO5);
        #[cfg(not(feature = "alt-log"))]
        let (tx, rx) = (esp_hal::gpio::NoPin, esp_hal::gpio::NoPin);

        Uart::new(peripherals.UART2, Default::default())
            .unwrap()
            .with_tx(tx)
            .with_rx(rx)
            .split()
    };

    #[cfg(feature = "usb-log")]
    logging::usb::setup_logging();
    #[cfg(feature = "alt-log")]
    logging::uart::setup_logging(uart_tx);
    #[cfg(feature = "rtt-log")]
    logging::rtt::setup_logging();

    // Set up allocators.
    ram::initialize(peripherals.PSRAM);

    // let mut io = Io::new(peripherals.IO_MUX);
    // io.set_interrupt_handler(interrupt_handler);

    // info!("IO Mux initialized!");

    // Enable pull-up on GPIO0 to prevent booting into download mode.
    let gpio0 = Output::new(
        peripherals.GPIO0,
        Level::High,
        OutputConfig::default().with_pull(Pull::Up),
    );

    // Enable LDO2
    let _ = Output::new(peripherals.GPIO17, Level::High, OutputConfig::default());

    // Enable antenna
    let _ = Output::new(peripherals.GPIO11, Level::Low, OutputConfig::default());

    let mut sha_backend = ShaBackend::new(peripherals.SHA);
    let _sha_driver_handle = sha_backend.start();

    let timg0 = TimerGroup::new(peripherals.TIMG0);
    let software_interrupt = SoftwareInterruptControl::new(peripherals.SW_INTERRUPT);
    esp_rtos::start(
        timg0.timer0,
        // software_interrupt.software_interrupt0,
    );

    // A task executor that is able to handle interrupts, and then return back to executing tasks.
    static EXECUTOR_CORE_0: StaticCell<InterruptExecutor<2>> = StaticCell::new();
    let executor_core_0 = InterruptExecutor::new(software_interrupt.software_interrupt2);
    let executor_core_0 = EXECUTOR_CORE_0.init(executor_core_0);
    let interrupt_core_0_spawner = executor_core_0.start(interrupt::Priority::Priority1);

    // static EXECUTOR_CORE_1: StaticCell<InterruptExecutor<3>> = StaticCell::new();
    // let executor_core_1 = InterruptExecutor::new(software_interrupt.software_interrupt3);
    // let executor_core_1 = EXECUTOR_CORE_1.init(executor_core_1);
    // let interrupt_core_1_spawner = executor_core_1.start(interrupt::Priority::Priority2);

    info!("ESP-RTOS started!");

    let main_task_peripherals = MainPeripherals {
        // high_priority_task_spawner: interrupt_core_1_spawner,
        uart_rx,
        software_interrupt1: software_interrupt.software_interrupt1,
        software_interrupt0: software_interrupt.software_interrupt0,
        RNG: peripherals.RNG,
        ADC1: peripherals.ADC1,
        USB0: peripherals.USB0,
        FLASH: peripherals.FLASH,
        DMA_CH0: peripherals.DMA_CH0,
        SPI2: peripherals.SPI2,
        CPU_CTRL: peripherals.CPU_CTRL,
        GPIO19: peripherals.GPIO19,
        GPIO20: peripherals.GPIO20,
        display: DisplayPeripherals {
            DMA_CH2: peripherals.DMA_CH2,
            LCD_CAM: peripherals.LCD_CAM,
            LEDC: peripherals.LEDC,
            GPIO0: gpio0,
            GPIO1: peripherals.GPIO1,
            GPIO2: peripherals.GPIO2,
            GPIO3: peripherals.GPIO3,
            GPIO4: peripherals.GPIO4,
            #[cfg(not(feature = "alt-log"))]
            GPIO5: peripherals.GPIO5,
            GPIO6: peripherals.GPIO6,
            #[cfg(not(feature = "alt-log"))]
            GPIO12: peripherals.GPIO12,
            GPIO13: peripherals.GPIO13,
            GPIO14: peripherals.GPIO14,
            GPIO15: peripherals.GPIO15,
            GPIO16: peripherals.GPIO16,
            GPIO21: peripherals.GPIO21,
            GPIO34: peripherals.GPIO34,
            GPIO35: peripherals.GPIO35,
            GPIO36: peripherals.GPIO36,
            GPIO37: peripherals.GPIO37,
            GPIO38: peripherals.GPIO38,
            GPIO39: peripherals.GPIO39,
            GPIO40: peripherals.GPIO40,
            GPIO41: peripherals.GPIO41,
            GPIO42: peripherals.GPIO42,
            GPIO43: peripherals.GPIO43,
            GPIO44: peripherals.GPIO44,
        },
        matrix: MatrixPeripherals {
            I2C0: peripherals.I2C0,
            GPIO7: peripherals.GPIO7,
            GPIO8: peripherals.GPIO8,
            GPIO9: peripherals.GPIO9,
        },
    };

    interrupt_core_0_spawner.must_spawn(main_task(main_task_peripherals));
}

/// Peripherals passed to the main task.
#[allow(non_snake_case)]
struct MainPeripherals {
    uart_rx: UartRx<'static, Blocking>,
    software_interrupt1: SoftwareInterrupt<'static, 1>,
    software_interrupt0: SoftwareInterrupt<'static, 0>,
    RNG: esp_hal::peripherals::RNG<'static>,
    ADC1: esp_hal::peripherals::ADC1<'static>,
    USB0: esp_hal::peripherals::USB0<'static>,
    FLASH: esp_hal::peripherals::FLASH<'static>,
    DMA_CH0: esp_hal::peripherals::DMA_CH0<'static>,
    SPI2: esp_hal::peripherals::SPI2<'static>,
    CPU_CTRL: esp_hal::peripherals::CPU_CTRL<'static>,
    // GPIO10: esp_hal::peripherals::GPIO10<'static>,
    // GPIO18: esp_hal::peripherals::GPIO18<'static>,
    GPIO19: esp_hal::peripherals::GPIO19<'static>,
    GPIO20: esp_hal::peripherals::GPIO20<'static>,
    // GPIO33: esp_hal::peripherals::GPIO33<'static>,
    // GPIO45: esp_hal::peripherals::GPIO45<'static>,
    // GPIO46: esp_hal::peripherals::GPIO46<'static>,
    // GPIO47: esp_hal::peripherals::GPIO47<'static>,
    // GPIO48: esp_hal::peripherals::GPIO48<'static>,
    display: DisplayPeripherals,
    matrix: MatrixPeripherals,
}

#[allow(non_snake_case)]
struct MatrixPeripherals {
    I2C0: esp_hal::peripherals::I2C0<'static>,
    GPIO7: esp_hal::peripherals::GPIO7<'static>,
    GPIO8: esp_hal::peripherals::GPIO8<'static>,
    GPIO9: esp_hal::peripherals::GPIO9<'static>,
}

#[embassy_executor::task]
async fn main_task(peripherals: MainPeripherals) {
    // let _spawner = unsafe { Spawner::for_current_executor() }.await;

    // Enable the TRNG source, so `Trng` can be constructed.
    let _trng_source = TrngSource::new(peripherals.RNG, peripherals.ADC1);

    #[cfg(feature = "ble")]
    let mut host_resources = rmk::HostResources::new();
    #[cfg(feature = "ble")]
    let stack = {
        use bt_hci::controller::ExternalController;
        use esp_radio::{Controller as RadioController, ble::controller::BleConnector};

        let mut rng = esp_hal::rng::Trng::try_new().unwrap();
        static RADIO: StaticCell<RadioController<'static>> = StaticCell::new();
        let radio = RADIO.init(esp_radio::init().unwrap());
        let bluetooth = peripherals.BT;
        let connector = BleConnector::new(radio, bluetooth, Default::default()).unwrap();
        let controller: ExternalController<_, 20> = ExternalController::new(connector);
        let central_addr = [0x18, 0xe2, 0x21, 0x80, 0xc0, 0xc7];
        let ble_stack =
            rmk::ble::build_ble_stack(controller, central_addr, &mut rng, &mut host_resources)
                .await;

        info!("BLE stack for RMK built!");

        ble_stack
    };

    // Initialize USB
    #[cfg(not(feature = "no-usb"))]
    let usb_driver = {
        use esp_hal::otg_fs::Usb;
        use esp_hal::otg_fs::asynch::{Config, Driver};

        static EP_MEMORY: StaticCell<[u8; 1024]> = StaticCell::new();
        let ep_memory = EP_MEMORY.init_with(|| [0_u8; _]);
        let usb = Usb::new(peripherals.USB0, peripherals.GPIO20, peripherals.GPIO19);
        // Create the driver, from the HAL.
        let config = Config::default();
        let driver = Driver::new(usb, ep_memory, config);

        info!("USB driver for RMK built!");

        driver
    };

    // Initialize the flash
    let flash_partitions = flash::initialize(peripherals.FLASH);

    info!("Flash memory configured!");

    // Uncomment this to run bounce buffer test code instead.
    // dpi::test_bounce_buffers(peripherals.DMA_CH0, peripherals.SPI2, peripherals.display).await;
    // return;

    // RMK config
    let vial_config = VialConfig::new(VIAL_KEYBOARD_ID, VIAL_KEYBOARD_DEF, &[(0, 0), (1, 1)]);
    let storage_config = StorageConfig {
        start_addr: 0,
        num_sectors: {
            assert!(
                flash_partitions.rmk.size() % FlashStorage::SECTOR_SIZE == 0,
                "The size of the RMK partition must be a multiple of {} bytes. Current size: {}",
                FlashStorage::SECTOR_SIZE,
                flash_partitions.rmk.size()
            );
            (flash_partitions.rmk.size() / FlashStorage::SECTOR_SIZE) as u8
        },
        ..Default::default()
    };
    // Retrieve the hardware-unique MAC address.
    let mac_address = Efuse::read_base_mac_address();
    static SERIAL_NUMBER: StaticCell<Box<str>> = StaticCell::new();
    let serial_number = SERIAL_NUMBER.init_with(|| {
        /// A magic prefix string that is required for the device to be recognized by the Vial GUI.
        const VIAL_SERIAL_PREFIX: &str = "vial:f64c2b3c";
        format!(
            "{VIAL_SERIAL_PREFIX}:acid:{:02x}",
            mac_address.iter().format("")
        )
        .into_boxed_str()
    });
    let device_config = DeviceConfig {
        vid: VIAL_VENDOR_ID,
        pid: VIAL_PRODUCT_ID,
        manufacturer: "",
        product_name: VIAL_KEYBOARD_NAME,
        serial_number,
    };
    info!("RMK Device Config: {device_config:#04x?}");
    let rmk_config = RmkConfig {
        device_config,
        vial_config,
        storage_config,
    };

    // Initialze keyboard stuffs
    // Initialize the storage and keymap
    let mut default_keymap = config::get_default_keymap();
    let mut behavior_config = config::get_behavior_config();
    let mut positional_config = config::get_positional_config();
    let (keymap, mut storage) = initialize_keymap_and_storage(
        &mut default_keymap,
        flash_partitions.rmk,
        &storage_config,
        &mut behavior_config,
        &mut positional_config,
    )
    .await;

    info!("Initialized keymap and storage for RMK!");

    let mut keyboard = Keyboard::new(&keymap); // Initialize the light controller

    info!("Keyboard initialized!");
    info!("Awaiting on all tasks...");

    // TODO: Probably want to select! instead and re-try.
    join_all![
        ram::run_alloc_stats_reporter(),
        initialize_and_run_rmk_devices(peripherals.matrix),
        keyboard.run(), // Keyboard is special
        run_rmk(
            &keymap,
            #[cfg(not(feature = "no-usb"))]
            usb_driver,
            #[cfg(feature = "ble")]
            &stack,
            &mut storage,
            rmk_config,
        ),
        create_hid_report_interceptor(),
        initialize_display_and_renderer(
            peripherals.display,
            peripherals.DMA_CH0,
            peripherals.SPI2,
            peripherals.CPU_CTRL,
            peripherals.software_interrupt0,
            peripherals.software_interrupt1,
            flash_partitions.acid
        ),
        console::run_console(peripherals.uart_rx.into_async())
    ]
    .await;
}

async fn initialize_and_run_rmk_devices(matrix_peripherals: MatrixPeripherals) {
    // Initialize the matrix and keyboard
    const I2C_ADDR_MATRIX_LEFT: I2cAddress = I2cAddress::SevenBit(0b0100000);
    const I2C_ADDR_MATRIX_RIGHT: I2cAddress = I2cAddress::SevenBit(0b0100001);

    let i2c = I2c::new(matrix_peripherals.I2C0, Default::default())
        .unwrap()
        .with_sda(matrix_peripherals.GPIO8)
        .with_scl(matrix_peripherals.GPIO9);
    let matrix_interrupt_low = Input::new(matrix_peripherals.GPIO7, InputConfig::default());
    let mut matrix = IoeMatrix::new(
        matrix_interrupt_low,
        i2c.into_async(),
        DefaultDebouncer::new(),
        [I2C_ADDR_MATRIX_LEFT, I2C_ADDR_MATRIX_RIGHT],
    )
    .await;
    run_devices! (
        (matrix) => rmk::channel::EVENT_CHANNEL,
    )
    .await
}

async fn initialize_display_and_renderer(
    display_peripherals: DisplayPeripherals,
    dma_ch0: esp_hal::peripherals::DMA_CH0<'static>,
    spi2: esp_hal::peripherals::SPI2<'static>,
    cpu_ctrl: esp_hal::peripherals::CPU_CTRL<'static>,
    software_interrupt0: SoftwareInterrupt<'static, 0>,
    software_interrupt1: SoftwareInterrupt<'static, 1>,
    partition_acid: Partition,
) {
    let st7701s = display_peripherals.into_display().await;

    static FRAMEBUFFER: StaticCell<Framebuffer> = StaticCell::new();
    let framebuffer = FRAMEBUFFER.init_with(move || {
        Framebuffer::new(
            dma_ch0,
            spi2.into(),
            st7701s.dpi,
            BurstConfig {
                internal_memory: InternalBurstConfig::Enabled,
                external_memory: ExternalBurstConfig::Size32,
            },
            // The burst config (16/32/64) doesn't seem to affect the alignment of the row size.
            //
            // |                      |     ( displayed range )     |
            // |                      [ pad ]                 [ pad ]
            // |                      [   DMA-transmissible range   ]
            // [ DMA-t. left overscan ]     |                 |     |
            // 0                     112   120               360   368 (index of u16 pixel)
            // ^ aligned              ^ aligned                     ^ aligned
            //
            // TODO: Compute the appropriate ranges to pass to the renderer and DPI peripheral.
            //       The renderer should pass the size of the `pad`ding to the GUI is parameters,
            //       to align the content to the displayed range.
            112,
            368 - 112,
            960,
            16,
            true,
        )
    });

    info!("Framebuffer created!");

    let window_size = [framebuffer.height, framebuffer.width];
    let swapchain_writer = framebuffer.swapchain.take().unwrap();

    static SECOND_CORE_STACK: StaticCell<Stack<{ STACK_SIZE_CORE_APP }>> = StaticCell::new();
    let second_core_stack = SECOND_CORE_STACK.init(Stack::new());
    esp_rtos::start_second_core(
        cpu_ctrl,
        software_interrupt0,
        software_interrupt1,
        second_core_stack,
        move || {
            // static EXECUTOR: StaticCell<InterruptExecutor<2>> = StaticCell::new();
            // let exec = EXECUTOR.init(InterruptExecutor::new(
            //     peripherals.software_interrupt2,
            // ));
            // let spawner = exec.start(Priority::Priority3);
            // spawner.must_spawn(run_renderer_task());
            static EXECUTOR: StaticCell<Executor> = StaticCell::new();
            let executor: &mut Executor = EXECUTOR.init(Executor::new());
            executor.run(|spawner| {
                let slint_backend = SlintBackend {
                    // peripherals: RefCell::new(Some(peripherals)),
                    window_size,
                    window: RefCell::new(None),
                    swapchain: RefCell::new(swapchain_writer),
                    quit_event_loop: Default::default(),
                    events: Arc::new(blocking_mutex::Mutex::new(RefCell::new(VecDeque::new()))),
                };
                spawner.must_spawn(ui::run_renderer_task(slint_backend, partition_acid));
            });
        },
    );

    info!("Second core started!");

    let bb_controller = DmaBounceController::new(framebuffer.bounce_buffers.take().unwrap());
    let mut user_controller = UserController::new(st7701s.controller, bb_controller);

    user_controller.event_loop().await
}

enum DmaBounceControllerInner {
    Idle(DmaBounce),
    Transmitting(RunningDmaBounceHandle),
}

struct DmaBounceController {
    inner: Option<DmaBounceControllerInner>,
}

impl DmaBounceController {
    pub fn new(dma_bounce: DmaBounce) -> Self {
        Self {
            inner: Some(DmaBounceControllerInner::Idle(dma_bounce)),
        }
    }

    pub async fn start(&mut self) -> Result<(), ()> {
        let DmaBounceControllerInner::Idle(dma_bounce) = self.inner.take().unwrap() else {
            return Err(());
        };
        self.inner = Some(DmaBounceControllerInner::Transmitting(
            dma_bounce.launch_interrupt_driven_task().await,
        ));
        Ok(())
    }

    pub async fn stop(&mut self) -> Result<(), ()> {
        let DmaBounceControllerInner::Transmitting(running_dma_bounce) = self.inner.take().unwrap()
        else {
            return Err(());
        };
        self.inner = Some(DmaBounceControllerInner::Idle(
            running_dma_bounce.stop().await,
        ));
        Ok(())
    }
}

struct UserController<'a> {
    sub: ControllerSub,
    lcd_controller: St7701sController<'a>,
    bb_controller: DmaBounceController,
}

impl<'a> UserController<'a> {
    fn new(lcd_controller: St7701sController<'a>, bb_controller: DmaBounceController) -> Self {
        Self {
            sub: CONTROLLER_CHANNEL.subscriber().unwrap(),
            lcd_controller,
            bb_controller,
        }
    }
}

impl<'a> Controller for UserController<'a> {
    type Event = ControllerEvent;

    async fn process_event(&mut self, event: Self::Event) {
        if let ControllerEvent::Key(keyboard_event, KeyAction::Single(Action::User(user_key_index))) =
            event
            && user_key_index == CustomKeycodes::FOCUS_LCD as u8
            && keyboard_event.pressed
        {
            let enabled = !LCD_ENABLED.fetch_xor(true, Ordering::SeqCst);

            match enabled {
                false => {
                    info!("Disabling LCD.");
                    *rmk::channel::KEYBOARD_REPORT_SENDER.write().await =
                        &rmk::channel::KEYBOARD_REPORT_RECEIVER;
                    self.lcd_controller.sleep_on().await;
                    self.bb_controller.stop().await.unwrap();
                }
                true => {
                    info!("Enabling LCD.");
                    *rmk::channel::KEYBOARD_REPORT_SENDER.write().await = &KEYBOARD_REPORT_PROXY;
                    self.bb_controller.start().await.unwrap();
                    self.lcd_controller.sleep_off().await;
                }
            }
        }
    }

    async fn next_message(&mut self) -> Self::Event {
        self.sub.next_message_pure().await
    }
}

// // TODO: Not needed currently. If it is ever enabled, don't forget to register it in Io.
// #[handler]
// #[ram] // Improves performance.
// fn interrupt_handler() {
//     // esp_println::println!(
//     //     "GPIO Interrupt with priority {}",
//     //     esp_hal::xtensa_lx::interrupt::get_level()
//     // );
// }
