#![no_std]
#![no_main]
#![feature(macro_metavar_expr)]

extern crate alloc;

mod keymap;
#[macro_use]
mod macros;
mod lcd;
mod matrix;
mod st7701s;
mod vial;

use core::alloc::Layout;
use core::ptr::addr_of_mut;

use alloc::boxed::Box;
use alloc::vec;
use bt_hci::controller::ExternalController;
use embassy_executor::Spawner;
use esp_alloc::{HeapRegion, MemoryCapability};
use esp_hal::clock::{CpuClock, RtcClock};
use esp_hal::delay::Delay;
use esp_hal::dma::{BurstConfig, DmaDescriptor, DmaTxBuf, ExternalBurstConfig};
use esp_hal::gpio::{Flex, Input, InputConfig, Level, Output, OutputConfig, Pull};
use esp_hal::i2c::master::{I2c, I2cAddress};
use esp_hal::lcd_cam::LcdCam;
use esp_hal::lcd_cam::lcd::dpi::{Dpi, Format, FrameTiming};
use esp_hal::lcd_cam::lcd::{ClockMode, Phase, Polarity};
use esp_hal::mcpwm::{McPwm, PeripheralClockConfig};
use esp_hal::otg_fs::Usb;
use esp_hal::otg_fs::asynch::{Config, Driver};
// use esp_hal::psram::{FlashFreq, PsramConfig, PsramSize, SpiRamFreq, SpiTimingConfigCoreClock};
use esp_hal::rng::TrngSource;
use esp_hal::rtc_cntl::sleep::RtcSleepConfig;
use esp_hal::rtc_cntl::{Rtc, sleep};
use esp_hal::time::Rate;
use esp_hal::timer::systimer::SystemTimer;
use esp_hal::timer::timg::TimerGroup;
use esp_radio::Controller;
use esp_radio::ble::controller::BleConnector;
use esp_storage::FlashStorage;
use log::{LevelFilter, info};
use rmk::ble::build_ble_stack;
use rmk::channel::EVENT_CHANNEL;
use rmk::config::{BehaviorConfig, PositionalConfig, RmkConfig, StorageConfig, VialConfig};
use rmk::debounce::default_debouncer::DefaultDebouncer;
use rmk::futures::future::join3;
use rmk::input_device::Runnable;
use rmk::keyboard::Keyboard;
use rmk::matrix::Matrix;
use rmk::storage::async_flash_wrapper;
use rmk::{HostResources, initialize_keymap_and_storage, run_devices, run_rmk};
use static_cell::StaticCell;
use {esp_alloc as _, esp_backtrace as _};

use crate::keymap::*;
use crate::lcd::spi_write;
use crate::matrix::IoeMatrix;
use crate::vial::{VIAL_KEYBOARD_DEF, VIAL_KEYBOARD_ID};

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

static PSRAM_ALLOCATOR: esp_alloc::EspHeap = esp_alloc::EspHeap::empty();

#[esp_rtos::main]
async fn main(spawner: Spawner) {
    esp_println::logger::init_logger_from_env();
    info!("Logger initialized!");

    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::default(),
    //     ram_frequency: SpiRamFreq::Freq80m,
    // });
    let peripherals: esp_hal::peripherals::Peripherals = esp_hal::init(config);
    info!("System initialized!");

    // Use the internal DRAM as the heap.
    esp_alloc::heap_allocator!(#[unsafe(link_section = ".dram2_uninit")] size: 64 * 1024);
    info!("Heap initialized! {:#?}", esp_alloc::HEAP.stats());

    // let timer0 = SystemTimer::new(peripherals.SYSTIMER);
    // esp_hal_embassy::init(timer0.alarm0);
    // info!("Embassy initialized!");

    // // Initialize the PSRAM allocator.
    // {
    //     let (psram_offset, psram_size) = esp_hal::psram::psram_raw_parts(&peripherals.PSRAM);
    //     unsafe {
    //         PSRAM_ALLOCATOR.add_region(HeapRegion::new(
    //             psram_offset,
    //             psram_size,
    //             MemoryCapability::External.into(),
    //         ));
    //     }
    // }
    // const BUFFER_LEN: usize = core::mem::size_of::<u16>()
    //     * (360 + /* TODO: Figure out why more bytes are needed: */ 8)
    //     * 960;
    // // Note: We just leak this buffer.
    // let buffer_ptr = unsafe {
    //     // ⚠️ Note: For chips that support DMA to/from PSRAM (ESP32-S3) DMA transfers to/from PSRAM
    //     // have extra alignment requirements. The address and size of the buffer pointed to by each
    //     // descriptor must be a multiple of the cache line (block) size. This is 32 bytes on
    //     // ESP32-S3.
    //     PSRAM_ALLOCATOR.alloc_caps(
    //         MemoryCapability::External.into(),
    //         Layout::from_size_align(BUFFER_LEN, 32).unwrap(),
    //     )
    // };
    // let buffer = unsafe { core::slice::from_raw_parts_mut(buffer_ptr, BUFFER_LEN) };
    // let burst_config: BurstConfig = ExternalBurstConfig::Size16.into();

    // info!(
    //     "PSRAM SPI burst config: max_compatible_chunk_size={}",
    //     burst_config.max_compatible_chunk_size()
    // );

    // let dma_buf_descs_len =
    //     esp_hal::dma::descriptor_count(BUFFER_LEN, burst_config.max_compatible_chunk_size(), false);
    // // Descriptors are initialized by `DmaTxBuf::new`.
    // let dma_buf_descs = Box::leak(vec![DmaDescriptor::EMPTY; dma_buf_descs_len].into_boxed_slice());
    // let mut dma_buf = DmaTxBuf::new(dma_buf_descs, buffer).unwrap();
    // // let mut dma_buf = dma_tx_buffer!(BUFFER_LEN).unwrap();

    // TODO: Spawn some tasks
    let _ = spawner;

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

    // TODO: Use PWM to control the pwm_pin.
    let mut _pwm = McPwm::new(peripherals.MCPWM0, PeripheralClockConfig::with_prescaler(1));
    let mut _pwm_pin = Output::new(peripherals.GPIO21, Level::High, OutputConfig::default());

    let mut matrix_interrupt = Input::new(
        peripherals.GPIO7,
        InputConfig::default().with_pull(Pull::Up),
    );

    // esp_alloc::heap_allocator!(size: 72 * 1024);
    let timg0 = TimerGroup::new(peripherals.TIMG0);
    #[cfg(target_arch = "riscv32")]
    let software_interrupt = SoftwareInterruptControl::new(peripherals.SW_INTERRUPT);
    esp_rtos::start(
        timg0.timer0,
        #[cfg(target_arch = "riscv32")]
        software_interrupt.software_interrupt0,
    );
    let _trng_source = TrngSource::new(peripherals.RNG, peripherals.ADC1);
    let mut rng = esp_hal::rng::Trng::try_new().unwrap();
    static RADIO: StaticCell<Controller<'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 mut host_resources = HostResources::new();
    let stack = build_ble_stack(controller, central_addr, &mut rng, &mut host_resources).await;

    // // Initialize USB
    #[cfg(not(feature = "no_usb"))]
    let usb_driver = {
        static mut EP_MEMORY: [u8; 1024] = [0; 1024];
        let usb = Usb::new(peripherals.USB0, peripherals.GPIO20, peripherals.GPIO19);
        // Create the driver, from the HAL.
        let config = Config::default();
        Driver::new(usb, unsafe { &mut *addr_of_mut!(EP_MEMORY) }, config)
    };

    // Initialize the flash
    let flash = FlashStorage::new(peripherals.FLASH);
    let flash = async_flash_wrapper(flash);

    let mut sck = Output::new(peripherals.GPIO36, Level::High, OutputConfig::default());
    let mut mosi = Flex::new(peripherals.GPIO35);
    let mut cs = Output::new(peripherals.GPIO6, Level::High, OutputConfig::default());

    mosi.set_input_enable(false);
    mosi.set_output_enable(true);

    // info!("init sequence writing");

    for (subsequence, delay_ms) in &*lcd::INIT_SEQUENCE_COMMANDS {
        for command in subsequence {
            spi_write(
                command.address(),
                command.args_iter().copied(),
                &mut mosi,
                &mut sck,
                &mut cs,
            );
            // info!("COMM 0x{:02X}", command.address());
            // for arg in command.args_iter() {
            //     info!("DATA 0x{arg:02X}");
            // }
        }
        Delay::new().delay_millis(*delay_ms as u32); // TODO: async?
        // Timer::after_millis(*delay_ms).await
    }

    // info!("init sequence written");

    let mut lcd = LcdCam::new(peripherals.LCD_CAM).into_async().lcd;
    let lcd_config = esp_hal::lcd_cam::lcd::dpi::Config::default()
        // Internal memory can use the full 16 MHz, but when external PSRAM is used, it cannot keep up with the display.
        // For that reason, we choose the highest value for which it doesn't glitch by showing black
        // stripes on the screen.
        //
        // There are three knobs you can turn to improve the bandwidth situation.
        // - increase psram frequency
        // - decrease the peripheral's frequency
        // - prevent flash from being used whilst your program is running. (There's a PR to make
        //   this easy to do)
        //   https://github.com/esp-rs/esp-hal/pull/3024
        .with_frequency(Rate::from_mhz(11)) // From Adafruit
        .with_clock_mode(ClockMode {
            polarity: Polarity::IdleLow, // From Adafruit
            phase: Phase::ShiftHigh,     // From Adafruit
        })
        .with_format(Format {
            enable_2byte_mode: true,
            ..Default::default()
        })
        .with_timing({
            // Adafruit's config for this LCD:
            // https://github.com/adafruit/Adafruit_CircuitPython_Qualia/blob/742d336e05e6a4d8bdaa46e15bbf60c9f30d2eba/adafruit_qualia/displays/bar240x960.py#L81-L97
            // https://github.com/adafruit/Adafruit_CircuitPython_Qualia/blob/742d336e05e6a4d8bdaa46e15bbf60c9f30d2eba/adafruit_qualia/displays/__init__.py#L59-L62
            // CircuitPython code handling Adafruit's config
            // https://github.com/adafruit/circuitpython/blob/97c6617817e95b1f6aa2ce458778aaa8371de39b/ports/espressif/common-hal/dotclockframebuffer/DotClockFramebuffer.c#L63
            // ESP-IDF peripheral configuration code:
            // https://github.com/espressif/esp-idf/blob/800f141f94c0f880c162de476512e183df671307/components/esp_lcd/rgb/esp_lcd_panel_rgb.c#L556
            // Espressif's docs:
            // https://docs.espressif.com/projects/esp-idf/en/v5.5.1/esp32s3/api-reference/peripherals/lcd/rgb_lcd.html#structures

            // TODO: Investigate PORCTRL instruction in datasheet of ST7701
            let horizontal_resolution: usize = 240;
            let vertical_resolution = 960;
            let overscan_left = 120;
            let vsync_width = 8;
            let hsync_width = 8;
            let horizontal_blank_front_porch = 20;
            let horizontal_blank_back_porch = 20;
            let vertical_blank_front_porch = 20;
            let vertical_blank_back_porch = 20;
            let hsync_position = 0;
            let horizontal_active_width = (horizontal_resolution + overscan_left).div_ceil(16) * 16; // Round up to a multiple of 16.
            let vertical_active_height = vertical_resolution;
            FrameTiming {
                horizontal_total_width: hsync_width
                    + horizontal_blank_back_porch
                    + horizontal_active_width
                    + horizontal_blank_front_porch,
                vertical_total_height: vsync_width
                    + vertical_blank_back_porch
                    + vertical_active_height
                    + vertical_blank_front_porch,
                horizontal_blank_front_porch: horizontal_blank_front_porch + hsync_width,
                vertical_blank_front_porch: vertical_blank_front_porch + vsync_width,
                horizontal_active_width,
                vertical_active_height,
                vsync_width,
                hsync_width,
                hsync_position,
            }
        })
        .with_hsync_idle_level(Level::High)
        .with_vsync_idle_level(Level::High)
        .with_de_idle_level(Level::Low);
    let mut dpi = Dpi::new(lcd, peripherals.DMA_CH2, lcd_config)
        .unwrap()
        .with_de(peripherals.GPIO37)
        .with_pclk(peripherals.GPIO34)
        .with_hsync(peripherals.GPIO44)
        .with_vsync(peripherals.GPIO43)
        // Blue
        .with_data0(peripherals.GPIO38)
        .with_data1(peripherals.GPIO39)
        .with_data2(peripherals.GPIO40)
        .with_data3(peripherals.GPIO41)
        .with_data4(peripherals.GPIO42)
        // Green
        .with_data5(peripherals.GPIO5)
        .with_data6(peripherals.GPIO12)
        .with_data7(peripherals.GPIO13)
        .with_data8(peripherals.GPIO14)
        .with_data9(peripherals.GPIO15)
        .with_data10(peripherals.GPIO16)
        // Red
        .with_data11(gpio0)
        .with_data12(peripherals.GPIO1)
        .with_data13(peripherals.GPIO2)
        .with_data14(peripherals.GPIO3)
        .with_data15(peripherals.GPIO4);

    // RMK config
    let vial_config = VialConfig::new(VIAL_KEYBOARD_ID, VIAL_KEYBOARD_DEF, &[(0, 0), (1, 1)]);
    let storage_config = StorageConfig {
        start_addr: 0x3f0000,
        num_sectors: 16,
        ..Default::default()
    };
    let rmk_config = RmkConfig {
        vial_config,
        storage_config,
        ..Default::default()
    };

    // Initialze keyboard stuffs
    // Initialize the storage and keymap
    let mut default_keymap = keymap::get_default_keymap();
    let mut behavior_config = BehaviorConfig::default();
    let mut per_key_config = PositionalConfig::default();
    let (keymap, mut storage) = initialize_keymap_and_storage(
        &mut default_keymap,
        flash,
        &storage_config,
        &mut behavior_config,
        &mut per_key_config,
    )
    .await;

    // 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(peripherals.I2C0, Default::default())
        .unwrap()
        .with_sda(peripherals.GPIO8)
        .with_scl(peripherals.GPIO9);

    let mut matrix = IoeMatrix::new(
        i2c.into_async(),
        DefaultDebouncer::new(),
        [I2C_ADDR_MATRIX_LEFT, I2C_ADDR_MATRIX_RIGHT],
    )
    .await;
    let mut keyboard = Keyboard::new(&keymap); // Initialize the light controller

    join3(
        run_devices! (
            (matrix) => EVENT_CHANNEL,
        ),
        keyboard.run(), // Keyboard is special
        run_rmk(
            &keymap,
            #[cfg(not(feature = "no_usb"))]
            usb_driver,
            &stack,
            &mut storage,
            rmk_config,
        ),
    )
    .await;
}

// #[esp_rtos::main]
// async fn main(_s: Spawner) {
//     // Initialize the peripherals and bluetooth controller
//     esp_println::logger::init_logger(LevelFilter::max());
//     let peripherals = esp_hal::init(esp_hal::Config::default().with_cpu_clock(CpuClock::max()));
//     esp_alloc::heap_allocator!(size: 72 * 1024);
//     let timg0 = TimerGroup::new(peripherals.TIMG0);
//     #[cfg(target_arch = "riscv32")]
//     let software_interrupt = SoftwareInterruptControl::new(peripherals.SW_INTERRUPT);
//     esp_rtos::start(
//         timg0.timer0,
//         #[cfg(target_arch = "riscv32")]
//         software_interrupt.software_interrupt0,
//     );
//     let _trng_source = TrngSource::new(peripherals.RNG, peripherals.ADC1);
//     let mut rng = esp_hal::rng::Trng::try_new().unwrap();
//     static RADIO: StaticCell<Controller<'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 mut host_resources = HostResources::new();
//     let stack = build_ble_stack(controller, central_addr, &mut rng, &mut host_resources).await;

//     // Initialize the flash
//     let flash = FlashStorage::new(peripherals.FLASH);
//     let flash = async_flash_wrapper(flash);

//     // Initialize the IO pins
//     let (row_pins, col_pins) = config_matrix_pins_esp!(peripherals: peripherals, input: [GPIO6, GPIO7, GPIO21, GPIO35], output: [GPIO3, GPIO4, GPIO5]);

//     // RMK config
//     let vial_config = VialConfig::new(VIAL_KEYBOARD_ID, VIAL_KEYBOARD_DEF, &[(0, 0), (1, 1)]);
//     let storage_config = StorageConfig {
//         start_addr: 0x3f0000,
//         num_sectors: 16,
//         ..Default::default()
//     };
//     let rmk_config = RmkConfig {
//         vial_config,
//         storage_config,
//         ..Default::default()
//     };

//     // Initialze keyboard stuffs
//     // Initialize the storage and keymap
//     let mut default_keymap = keymap::get_default_keymap();
//     let mut behavior_config = BehaviorConfig::default();
//     let mut per_key_config = PositionalConfig::default();
//     let (keymap, mut storage) = initialize_keymap_and_storage(
//         &mut default_keymap,
//         flash,
//         &storage_config,
//         &mut behavior_config,
//         &mut per_key_config,
//     )
//     .await;

//     // Initialize the matrix and keyboard
//     let debouncer = DefaultDebouncer::new();
//     let mut i2c = I2c::new(
//         peripherals.I2C0,
//         esp_hal::i2c::master::Config::default().with_frequency(Rate::from_khz(400)),
//     )
//     .unwrap()
//     .with_sda(peripherals.GPIO8)
//     .with_scl(peripherals.GPIO9);

//     const I2C_ADDR_MATRIX_LEFT: I2cAddress = I2cAddress::SevenBit(0b0100000);
//     const I2C_ADDR_MATRIX_RIGHT: I2cAddress = I2cAddress::SevenBit(0b0100001);

//     let mut matrix = IoeMatrix::new(
//         i2c.into_async(),
//         debouncer,
//         [I2C_ADDR_MATRIX_LEFT, I2C_ADDR_MATRIX_RIGHT],
//     )
//     .await;
//     // let mut matrix = Matrix::<_, _, _, ROW, COL, true>::new(row_pins, col_pins, debouncer);
//     // let mut matrix = rmk::matrix::TestMatrix::<ROW, COL>::new();
//     let mut keyboard = Keyboard::new(&keymap); // Initialize the light controller

//     join3(
//         run_devices! (
//             (matrix) => EVENT_CHANNEL,
//         ),
//         keyboard.run(), // Keyboard is special
//         run_rmk(&keymap, &stack, &mut storage, rmk_config),
//     )
//     .await;
// }