#![no_std]
#![no_main]

//! # Pinout
//!
//! GP0..GP14 => A0..A14
//! GP15..GP22 => D0..D7
//! GP26 => CE
//! GP27 => OE

mod serial;

use core::{
    mem::MaybeUninit,
    sync::atomic::{AtomicU8, Ordering},
    task::Poll,
};

use defmt_rtt as _;
use embassy_executor::{InterruptExecutor, Spawner, task};
use embassy_futures::{poll_once, yield_now};
use embassy_rp::{
    bind_interrupts,
    clocks::ClockConfig,
    flash::{self, FLASH_BASE, Flash},
    gpio::{Drive, Level, Output, SlewRate},
    interrupt::{self, InterruptExt, Priority},
    peripherals::{FLASH, PIO0, USB},
    pio::{self, Direction, Pio, ShiftConfig, ShiftDirection, StateMachine, program::pio_asm},
    usb::{self, Driver},
};
use embassy_usb::{
    UsbDevice,
    class::cdc_acm::{self, CdcAcmClass},
};
use panic_probe as _;
use static_cell::{ConstStaticCell, StaticCell};

use crate::serial::{Link, PacketBuilder};

const ADDRESS_PINS: u8 = 15;
const DATA_PINS: u8 = 8;

const ROM_SIZE: usize = 32 * 1024;
const FLASH_SIZE: usize = 2 * 1024 * 1024;

static ROM_DATA: [AtomicU8; ROM_SIZE] = [const { AtomicU8::new(0) }; ROM_SIZE];

#[unsafe(link_section = ".romdata")]
#[used]
static INIT_ROM_DATA: MaybeUninit<[u8; ROM_SIZE]> = MaybeUninit::uninit();

bind_interrupts!(struct Irqs {
    PIO0_IRQ_0 => pio::InterruptHandler<PIO0>;
    USBCTRL_IRQ => usb::InterruptHandler<USB>;
});

static EXECUTOR_HIGH: InterruptExecutor = InterruptExecutor::new();

#[embassy_rp::interrupt]
unsafe fn SWI_IRQ_0() {
    unsafe { EXECUTOR_HIGH.on_interrupt() }
}

#[embassy_executor::main]
async fn main(spawner: Spawner) -> ! {
    let config =
        embassy_rp::config::Config::new(defmt::unwrap!(ClockConfig::system_freq(133_000_000)));
    let p = embassy_rp::init(config);

    let led = unsafe { p.PIN_25.clone_unchecked() };

    let mut flash = Flash::<_, _, FLASH_SIZE>::new_blocking(p.FLASH);

    if let Err(e) = load_rom(&mut flash) {
        defmt::error!("Unable to initialize rom from flash: {}", e);
    }

    // Pull from fifo and write to all data pins
    let data = pio_asm!(".wrap_target", "out pins, 8", ".wrap");
    // OR CE and OE and set 4 pindirs
    let output_enable = pio_asm!(
        ".origin 0",
        ".side_set 4 pindirs",
        "mov pc, pins side 0xf",
        "mov pc, pins side 0x0",
        "mov pc, pins side 0x0",
        "mov pc, pins side 0x0"
    );
    // report pin activity to the cpu
    let report_data = pio_asm!(
        "wait_oe:",
        "mov y, pins",      // read all pins
        "mov osr, y",       // move into the osr so we can shift them out
        "out null, 26",     // skipp all pins until CE
        "out x, 2",         // put CE and OE into x
        "jmp x--, wait_oe", // if any of CE and OE are set go back to the start
        "mov isr, y",       // put the address into the isr
        "push",
        ".wrap_target",
        "mov x, pins",      // read all pins
        "jmp x!=y wait_oe", // if any pins have changed, restart and check oe again
        ".wrap"
    );

    let Pio {
        mut common,
        mut sm0,
        mut sm1,
        mut sm2,
        mut sm3,
        ..
    } = Pio::new(p.PIO0, Irqs);
    let output_enable = common.load_program(&output_enable.program);

    let mut address_pins = [
        common.make_pio_pin(p.PIN_0),
        common.make_pio_pin(p.PIN_1),
        common.make_pio_pin(p.PIN_2),
        common.make_pio_pin(p.PIN_3),
        common.make_pio_pin(p.PIN_4),
        common.make_pio_pin(p.PIN_5),
        common.make_pio_pin(p.PIN_6),
        common.make_pio_pin(p.PIN_7),
        common.make_pio_pin(p.PIN_8),
        common.make_pio_pin(p.PIN_9),
        common.make_pio_pin(p.PIN_10),
        common.make_pio_pin(p.PIN_11),
        common.make_pio_pin(p.PIN_12),
        common.make_pio_pin(p.PIN_13),
        common.make_pio_pin(p.PIN_14),
    ];
    for pin in &mut address_pins {
        pin.set_schmitt(true);
    }
    let mut data_pins = [
        common.make_pio_pin(p.PIN_15),
        common.make_pio_pin(p.PIN_16),
        common.make_pio_pin(p.PIN_17),
        common.make_pio_pin(p.PIN_18),
        common.make_pio_pin(p.PIN_19),
        common.make_pio_pin(p.PIN_20),
        common.make_pio_pin(p.PIN_21),
        common.make_pio_pin(p.PIN_22),
    ];
    for pin in &mut data_pins {
        pin.set_slew_rate(SlewRate::Fast);
        pin.set_drive_strength(Drive::_12mA);
        pin.set_schmitt(true);
    }
    let data_pins = [
        &data_pins[0],
        &data_pins[1],
        &data_pins[2],
        &data_pins[3],
        &data_pins[4],
        &data_pins[5],
        &data_pins[6],
        &data_pins[7],
    ];
    let mut ce_pin = common.make_pio_pin(p.PIN_26);
    ce_pin.set_schmitt(true);
    let mut oe_pin = common.make_pio_pin(p.PIN_27);
    oe_pin.set_schmitt(true);
    let led_pin = common.make_pio_pin(p.PIN_25);
    let address_pins = [
        &address_pins[0],
        &address_pins[1],
        &address_pins[2],
        &address_pins[3],
        &address_pins[4],
        &address_pins[5],
        &address_pins[6],
        &address_pins[7],
        &address_pins[8],
        &address_pins[9],
        &address_pins[10],
        &address_pins[11],
        &address_pins[12],
        &address_pins[13],
        &address_pins[14],
        data_pins[0],
        data_pins[1],
        data_pins[2],
        data_pins[3],
        data_pins[4],
        data_pins[5],
        data_pins[6],
        data_pins[7],
        &common.make_pio_pin(p.PIN_23),
        &common.make_pio_pin(p.PIN_24),
        &led_pin,
        &ce_pin,
        &oe_pin,
    ];

    // configure data pio
    let mut cfg = pio::Config::default();
    cfg.use_program(&common.load_program(&data.program), &[]);
    cfg.set_out_pins(&data_pins[..]);
    cfg.shift_out = ShiftConfig {
        threshold: DATA_PINS,
        direction: ShiftDirection::default(),
        auto_fill: true,
    };
    sm0.set_config(&cfg);
    sm0.set_enable(true);

    // configure output enable pio
    let mut cfg = pio::Config::default();
    cfg.use_program(&output_enable, &data_pins[..4]);
    cfg.set_in_pins(&[&ce_pin, &oe_pin]);
    sm1.set_config(&cfg);
    sm1.set_enable(true);
    cfg.use_program(&output_enable, &data_pins[4..]);
    sm2.set_config(&cfg);
    sm2.set_enable(true);

    // configure address pio
    let mut cfg = pio::Config::default();
    cfg.use_program(&common.load_program(&report_data.program), &[]);
    cfg.set_in_pins(&address_pins[..]);
    cfg.shift_in = ShiftConfig {
        threshold: ADDRESS_PINS,
        direction: ShiftDirection::default(),
        auto_fill: true,
    };
    sm3.set_pin_dirs(Direction::In, &address_pins[..]);
    sm3.set_pin_dirs(Direction::Out, &[&led_pin]);
    sm3.set_config(&cfg);
    sm3.set_enable(true);

    let driver = Driver::new(p.USB, Irqs);

    let config = {
        let mut config = embassy_usb::Config::new(0x2e8a, 0x000a);
        config.manufacturer = Some("Max Känner");
        config.product = Some("PicoROM.rs");
        config.serial_number = Some("2");
        config.max_power = 100;
        config.max_packet_size_0 = 64;
        config
    };

    let mut builder = {
        static CONFIG_DESCRIPTOR: ConstStaticCell<[u8; 256]> = ConstStaticCell::new([0; 256]);
        static BOS_DESCRIPTOR: ConstStaticCell<[u8; 256]> = ConstStaticCell::new([0; 256]);
        static CONTROL_BUF: ConstStaticCell<[u8; 64]> = ConstStaticCell::new([0; 64]);
        embassy_usb::Builder::new(
            driver,
            config,
            CONFIG_DESCRIPTOR.take(),
            BOS_DESCRIPTOR.take(),
            &mut [],
            CONTROL_BUF.take(),
        )
    };

    let mut class = {
        static STATE: StaticCell<cdc_acm::State> = StaticCell::new();
        CdcAcmClass::new(&mut builder, STATE.init(cdc_acm::State::new()), 64)
    };

    let usb = builder.build();

    interrupt::PIO0_IRQ_0.set_priority(Priority::P0);
    interrupt::SWI_IRQ_0.set_priority(Priority::P1);
    interrupt::USBCTRL_IRQ.set_priority(Priority::P3);

    let int_spawner = EXECUTOR_HIGH.start(interrupt::SWI_IRQ_0);
    int_spawner.must_spawn(pio_task(sm0, sm3));

    spawner.must_spawn(usb_task(usb));

    let mut led = Output::new(led, Level::Low);

    loop {
        'outer: {
            led.set_low();
            class.wait_connection().await;
            while !class.rts() {
                yield_now().await;
            }
            defmt::info!("USB Connected");
            led.set_high();
            if let Err(e) = class.write_packet(b"PicoROM Hello").await {
                defmt::error!("Error sending preamble: {}", e);
                break 'outer;
            }
            let mut buf = [0; 64];
            let mut packet_builder = PacketBuilder::new();
            let mut link = Link::new();
            while class.rts() {
                let n = match poll_once(class.read_packet(&mut buf)) {
                    Poll::Ready(Ok(n)) => n,
                    Poll::Ready(Err(e)) => {
                        defmt::error!("Error reading packet data: {}", e);
                        break 'outer;
                    }
                    Poll::Pending => {
                        yield_now().await;
                        continue;
                    }
                };
                for byte in buf.into_iter().take(n) {
                    if let Some(packet) = packet_builder.push(byte) {
                        defmt::debug!("Got packet: {}", packet);
                        if let Some(response) = link.handle_packet(&packet, &mut flash) {
                            defmt::info!("Sending Response: {}", response);
                            if let Err(e) = response.send(&mut class).await {
                                defmt::error!("Unable to send response: {}", e);
                                break 'outer;
                            }
                        }
                    }
                }
            }
        }
        defmt::info!("USB Disconnected");
    }
}

#[task]
async fn usb_task(mut usb: UsbDevice<'static, Driver<'static, USB>>) -> ! {
    usb.run().await
}

#[task]
async fn pio_task(
    mut data_sm: StateMachine<'static, PIO0, 0>,
    mut address_sm: StateMachine<'static, PIO0, 3>,
) -> ! {
    loop {
        let address = address_sm.rx().wait_pull().await;
        let address = (address & 0x7fff) as u16;
        let data = ROM_DATA[usize::from(address)].load(Ordering::SeqCst);
        defmt::trace!("replying with {:#04x} @ {:#06x}", data, address);
        data_sm.tx().try_push(u32::from(data));
    }
}

fn store_rom(
    flash: &mut Flash<'_, FLASH, impl flash::Mode, FLASH_SIZE>,
) -> Result<(), flash::Error> {
    let offset = unsafe {
        u32::try_from(
            INIT_ROM_DATA
                .as_ptr()
                .cast::<u8>()
                .offset_from_unsigned(FLASH_BASE.cast::<u8>()),
        )
        .map_err(|_| flash::Error::OutOfBounds)?
    };
    let len = size_of_val(&INIT_ROM_DATA).min(ROM_DATA.len());
    defmt::info!("Erasing flash at offset {:#x} with size {:#x}", offset, len);
    flash.blocking_erase(
        offset,
        offset + u32::try_from(len).map_err(|_| flash::Error::OutOfBounds)?,
    )?;
    let rom_buffer = unsafe { &*(&raw const ROM_DATA[..len] as *const [u8]) };
    defmt::info!("Programming flash with buffer at {}", rom_buffer.as_ptr());
    flash.blocking_write(offset, rom_buffer)?;
    defmt::info!("Successfully commited rom to flash");
    Ok(())
}

fn load_rom(
    flash: &mut Flash<'_, FLASH, impl flash::Mode, FLASH_SIZE>,
) -> Result<(), flash::Error> {
    let offset = unsafe {
        u32::try_from(
            INIT_ROM_DATA
                .as_ptr()
                .cast::<u8>()
                .offset_from_unsigned(FLASH_BASE.cast::<u8>()),
        )
        .map_err(|_| flash::Error::OutOfBounds)?
    };
    let len = size_of_val(&INIT_ROM_DATA).min(ROM_DATA.len());
    for (i, rom) in ROM_DATA.iter().enumerate().take(len) {
        let mut init = [0u8];
        flash.blocking_read(
            offset + u32::try_from(i).map_err(|_| flash::Error::OutOfBounds)?,
            &mut init,
        )?;
        let init = init[0];
        rom.store(init, Ordering::SeqCst);
    }
    Ok(())
}