#![no_std]
#![no_main]

use rp_pico::entry;

use defmt as _;
use defmt_rtt as _;
use panic_probe as _;

use embedded_hal::digital::StatefulOutputPin;

use rp_pico::hal::pac;

use rp_pico::hal;

use usb_device::{class_prelude::*, prelude::*};

use usbd_storage::subclass::scsi::{Scsi, ScsiCommand};
use usbd_storage::subclass::Command;
use usbd_storage::transport::bbb::{BulkOnly, BulkOnlyError};
use usbd_storage::transport::TransportError;

const DISK_BLOCK_SIZE: u32 = 512;
const DISK_BLOCK_NUM: u32 = 256;

const USB_PACKET_SIZE: u16 = 64; // 8,16,32,64
const MAX_LUN: u8 = 0; // max 0x0F

const INQUIRY_RESPONSE: [u8; 36] = {
    let vendor: [u8; 8] = *b"jalr\0\0\0\0";
    let model: [u8; 16] = *b"USB RAM disk\0\0\0\0";
    let rev: [u8; 4] = *b"1312";
    [
        0x00, // periph qualifier, periph device type
        0x80, // Removable
        0x04, // SPC-2 compliance
        0x02, // NormACA, HiSu, Response data format
        0x20, // 36 bytes in total
        0x00, // additional fields, none set
        0x00, // additional fields, none set
        0x00, // additional fields, none set
        vendor[0], vendor[1],  vendor[2], vendor[3],
        vendor[4], vendor[5],  vendor[6], vendor[7],
        model[0],   model[1],   model[2],  model[3],
        model[4],   model[5],   model[6],  model[7],
        model[8],   model[9],  model[10], model[11],
        model[12], model[13],  model[14], model[15],
        rev[0],       rev[1],     rev[2],    rev[3]
    ]
};

#[derive(Default)]
struct State {
    storage_offset: usize,
    sense_key: Option<u8>,
    sense_key_code: Option<u8>,
    sense_qualifier: Option<u8>,
}

impl State {
    fn reset(&mut self) {
        self.storage_offset = 0;
        self.sense_key = None;
        self.sense_key_code = None;
        self.sense_qualifier = None;
    }
}

#[entry]
fn main() -> ! {
    let mut pac = pac::Peripherals::take().unwrap();

    let mut watchdog = hal::Watchdog::new(pac.WATCHDOG);

    let clocks = hal::clocks::init_clocks_and_plls(
        rp_pico::XOSC_CRYSTAL_FREQ,
        pac.XOSC,
        pac.CLOCKS,
        pac.PLL_SYS,
        pac.PLL_USB,
        &mut pac.RESETS,
        &mut watchdog,
    )
    .ok()
    .unwrap();

    let timer = hal::timer::Timer::new(pac.TIMER, &mut pac.RESETS, &clocks);

    let sio = hal::Sio::new(pac.SIO);
    let pins = rp_pico::Pins::new(
        pac.IO_BANK0,
        pac.PADS_BANK0,
        sio.gpio_bank0,
        &mut pac.RESETS,
    );

    let mut led_pin = pins.led.into_push_pull_output();

    let usb_bus = UsbBusAllocator::new(hal::usb::UsbBus::new(
        pac.USBCTRL_REGS,
        pac.USBCTRL_DPRAM,
        clocks.usb_clock,
        true,
        &mut pac.RESETS,
    ));

    let mut usb_transport_buf: [u8; 512] = [0; 512];
    let mut scsi = usbd_storage::subclass::scsi::Scsi::new(&usb_bus, USB_PACKET_SIZE, MAX_LUN, usb_transport_buf.as_mut_slice()).unwrap();

    let mut usb_device = UsbDeviceBuilder::new(&usb_bus, UsbVidPid(0xcafe, 0x4002))
        .strings(&[StringDescriptors::new(LangID::EN)
            .manufacturer("jalr")
            .product("RAM Mass Storage")])
        .unwrap()
        .device_class(0x08)
        .self_powered(false)
        .build();

    let mut storage: [u8; (DISK_BLOCK_SIZE * DISK_BLOCK_NUM) as usize] = [0u8; (DISK_BLOCK_SIZE * DISK_BLOCK_NUM) as usize];

    const LBA: u32 = DISK_BLOCK_NUM - 1;
    let (last_cyl, last_head, last_sect) = {
        const HPC: u32 = 16;
        const SPT: u32 = 63;
        (
            LBA / (HPC * SPT),
            (LBA / SPT) % HPC,
            (LBA % SPT) + 1
        )
    };

    const OEM: [u8; 8] = *b"mkfs.fat";
    const FILESYSTEM_ID: [u8; 4] = [0xDE, 0xAD, 0xBE, 0xEF];
    const FILESYSTEM_LABEL: [u8; 11] = *b"RAM_USB    ";

    storage[0x1B9] = 0x10;
    storage[0x1BE..0x1C3].copy_from_slice(&[0x00, 0x00, 0x02, 0x00, 0x0c]);
    storage[0x1C3] = last_head as u8;
    storage[0x1C4] = ((last_cyl >> 8) & 0xc0) as u8 | ((last_sect as u8) & 0x3f);
    storage[0x1C5] = last_cyl as u8;
    storage[0x1C6..0x1CA].copy_from_slice(&[0x01, 0x00, 0x00, 0x00]); // LBA of first absolute sector in the partition
    storage[0x1CA] = LBA as u8;
    storage[0x1CB] = (LBA >> 8) as u8;
    storage[0x1CC] = (LBA >> 16) as u8;
    storage[0x1CD] = (LBA >> 24) as u8;
    storage[0x1FE..0x200].copy_from_slice(&[0x55, 0xaa]); // boot signature

    // FAT12 Partition
    storage[0x200..0x203].copy_from_slice(&[0xEB, 0x3C, 0x90]);
    storage[0x203..0x20B].copy_from_slice(&OEM);
    storage[0x20B..0x21d].copy_from_slice(&[0x00, 0x02, 0x04, 0x01, 0x00,
        0x02, 0x00, 0x02, (LBA & 0xff) as u8,
                                (LBA >> 8) as u8,
                                      0xF8, 0x01, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01]);
    storage[0x224..0x227].copy_from_slice(&[0x80, 0x00, 0x29]);
    storage[0x227..0x22B].copy_from_slice(&FILESYSTEM_ID);
    storage[0x22B..0x236].copy_from_slice(&FILESYSTEM_LABEL);
    storage[0x236..0x23E].copy_from_slice(b"FAT12   ");
    storage[0x3FE..0x400].copy_from_slice(&[0x55, 0xAA]);

    storage[0x400..0x403].copy_from_slice(&[0xF8, 0xFF, 0xFF]);
    storage[0x600..0x603].copy_from_slice(&[0xF8, 0xFF, 0xFF]);
    storage[0x800..0x80B].copy_from_slice(&FILESYSTEM_LABEL);
    storage[0x80B] = 0x08; // file attribute 0x08 = volume label

    // https://de.wikipedia.org/wiki/File_Allocation_Table
    // Remove bootloader code: dd if=/dev/zero conv=notrunc of=foo bs=1 count=420 seek=602

    let mut state: State = State {
        storage_offset: 0,
        sense_key: None,
        sense_key_code: None,
        sense_qualifier: None,
    };

    defmt::info!("entering main loop");

    let mut start_time: u64 = 0;
    loop {
        let interval = if timer.get_counter().ticks() < 2_000_000 {
            25_000
        }
        else {
            match (usb_device.state(), led_pin.is_set_high().unwrap_or_else(|_| false)) {
                // The USB device has just been created or reset.
                (UsbDeviceState::Default, true) => 100_000,
                (UsbDeviceState::Default, false) => 100_000,
                // The USB device has received an address from the host.
                (UsbDeviceState::Addressed, true) => 250_000,
                (UsbDeviceState::Addressed, false) => 250_000,
                // The USB device has been configured and is fully functional.
                (UsbDeviceState::Configured, true) => 1_000_000,
                (UsbDeviceState::Configured, false) => 0,
                // The USB device has been suspended by the host or it has been unplugged from the USB bus.
                (UsbDeviceState::Suspend, true) => 25_000,
                (UsbDeviceState::Suspend, false) => 2_000_000,
            }
        };
        if timer.get_counter().ticks() - start_time > interval {
            let _ = led_pin.toggle();
            start_time = timer.get_counter().ticks();
        }

        if !usb_device.poll(&mut [&mut scsi]) {
            continue;
        }

        if matches!(usb_device.state(), UsbDeviceState::Default) {
            state.reset();
        }

        let _ = scsi.poll(|command| {
            if let Err(err) = process_command(command, &mut state, &mut storage) {
                defmt::error!("Got error: {}", err);
            }
        });
    }
}

fn process_command(
    mut command: Command<ScsiCommand, Scsi<BulkOnly<hal::usb::UsbBus, &mut [u8]>>>,
    state: &mut State,
    storage: &mut [u8; (DISK_BLOCK_SIZE * DISK_BLOCK_NUM) as usize]
) -> Result<(), TransportError<BulkOnlyError>> {
    defmt::info!("Handling: {}", command.kind);

    match command.kind {
        ScsiCommand::TestUnitReady { .. } => {
            command.pass();
        }
        ScsiCommand::Inquiry { .. } => {
            command.try_write_data_all(&INQUIRY_RESPONSE)?;
            command.pass();
        }
        ScsiCommand::RequestSense { .. } => {
            command.try_write_data_all(&[
                0x70,                         // RESPONSE CODE. Set to 70h for information on current errors
                0x00,                         // obsolete
                state.sense_key.unwrap_or(0), // Bits 3..0: SENSE KEY. Contains information describing the error.
                0x00,
                0x00,
                0x00,
                0x00, // INFORMATION. Device-specific or command-specific information.
                0x00, // ADDITIONAL SENSE LENGTH.
                0x00,
                0x00,
                0x00,
                0x00,                               // COMMAND-SPECIFIC INFORMATION
                state.sense_key_code.unwrap_or(0),  // ASC
                state.sense_qualifier.unwrap_or(0), // ASCQ
                0x00,
                0x00,
                0x00,
                0x00,
            ])?;
            state.reset();
            command.pass();
        },
        ScsiCommand::ReadCapacity10 { .. } => {
            let mut data = [0u8; 8];
            let _ = &mut data[0..4].copy_from_slice(&u32::to_be_bytes(DISK_BLOCK_NUM - 1));
            let _ = &mut data[4..8].copy_from_slice(&u32::to_be_bytes(DISK_BLOCK_SIZE));
            command.try_write_data_all(&data)?;
            command.pass();
        }
        ScsiCommand::ReadCapacity16 { .. } => {
            let mut data = [0u8; 16];
            let _ = &mut data[0..8].copy_from_slice(&u32::to_be_bytes(DISK_BLOCK_NUM - 1));
            let _ = &mut data[8..12].copy_from_slice(&u32::to_be_bytes(DISK_BLOCK_SIZE));
            command.try_write_data_all(&data)?;
            command.pass();
        }
        ScsiCommand::ReadFormatCapacities { .. } => {
            let mut data = [0u8; 12];
            let _ = &mut data[0..4].copy_from_slice(&[
                0x00, 0x00, 0x00, 0x08, // capacity list length
            ]);
            let _ = &mut data[4..8].copy_from_slice(&u32::to_be_bytes(DISK_BLOCK_NUM as u32)); // number of blocks
            data[8] = 0x01; //unformatted media
            let block_length_be = u32::to_be_bytes(DISK_BLOCK_SIZE);
            data[9] = block_length_be[1];
            data[10] = block_length_be[2];
            data[11] = block_length_be[3];

            command.try_write_data_all(&data)?;
            command.pass();
        }
        ScsiCommand::Read { lba, len } => {
            let lba = lba as u32;
            let len = len as u32;
            if state.storage_offset != (len * DISK_BLOCK_SIZE) as usize {
                let start = (DISK_BLOCK_SIZE * lba) as usize + state.storage_offset;
                let end = (DISK_BLOCK_SIZE * lba) as usize + (DISK_BLOCK_SIZE * len) as usize;

                // Uncomment this in order to push data in chunks smaller than a USB packet.
                // let end = min(start + USB_PACKET_SIZE as usize - 1, end);

                defmt::info!("Data transfer >>>>>>>> [{}..{}]", start, end);
                let count = command.write_data(&mut storage[start..end])?;
                state.storage_offset += count;
            } else {
                command.pass();
                state.storage_offset = 0;
            }
        },
        ScsiCommand::Write { lba, len } => {
            let lba = lba as u32;
            let len = len as u32;
            if state.storage_offset != (len * DISK_BLOCK_SIZE) as usize {
                let start = (DISK_BLOCK_SIZE * lba) as usize + state.storage_offset;
                let end = (DISK_BLOCK_SIZE * lba) as usize + (DISK_BLOCK_SIZE * len) as usize;
                defmt::info!("Data transfer <<<<<<<< [{}..{}]", start, end);
                let count = command.read_data(&mut storage[start..end])?;
                state.storage_offset += count;

                if state.storage_offset == (len * DISK_BLOCK_SIZE) as usize {
                    command.pass();
                    state.storage_offset = 0;
                }
            } else {
                command.pass();
                state.storage_offset = 0;
            }
        },
        ScsiCommand::ModeSense6 { .. } => {
            command.try_write_data_all(&[
                0x03, // number of bytes that follow
                0x00, // the media type is SBC
                0x00, // not write-protected, no cache-control bytes support
                0x00, // no mode-parameter block descriptors
            ])?;
            command.pass();
        }
        ScsiCommand::ModeSense10 { .. } => {
            command.try_write_data_all(&[0x00, 0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])?;
            command.pass();
        }
        ref unknown_scsi_kind => {
            defmt::error!("Unknown SCSI command: {}", unknown_scsi_kind);
            state.sense_key.replace(0x05); // illegal request Sense Key
            state.sense_key_code.replace(0x20); // Invalid command operation ASC
            state.sense_qualifier.replace(0x00); // Invalid command operation ASCQ
            command.fail();
        }
    }

    Ok(())
}