//! Route stdout through a USB serial interface
//!
//! You'll notice this looks pretty similar to imxrt-log's usbd backend.
//! I replace the bbqueue with a fusible Queue and BlockPool, and I run
//! the "backend" on a fusible Thread.
//!
//! You're responsible for registering the USB interrupt and calling
//! `on_interrupt` in that interrupt. You'll also need to call `start`
//! to set up the driver.
//!
//! This only supports stdout; I haven't gotten stdin working yet.

#![no_std]
#![feature(rustc_private)]
extern crate fusible;

use usb_device::device::UsbDeviceState;

const VID_PID: usb_device::device::UsbVidPid = usb_device::device::UsbVidPid(0x5824, 0x27dd);
const PRODUCT: &str = "rust-threadx-usbd";

/// Provide some extra overhead for the interrupt endpoint.
///
/// If you start noticing panics, check to make sure that this buffer
/// is large enough for all the max packet sizes for all the endpoints.
const ENDPOINT_BYTES: usize = MAX_PACKET_SIZE * 2 + EP0_CONTROL_PACKET_SIZE * 2 + 128;
static ENDPOINT_MEMORY: imxrt_usbd::EndpointMemory<ENDPOINT_BYTES> =
    imxrt_usbd::EndpointMemory::new();
static ENDPOINT_STATE: imxrt_usbd::EndpointState<6> = imxrt_usbd::EndpointState::new();

type Bus = imxrt_usbd::BusAdapter;
type Class<'a> = usbd_serial::CdcAcmClass<'a, Bus>;
type Device<'a> = usb_device::device::UsbDevice<'a, Bus>;

/// High-speed bulk endpoint limit.
const MAX_PACKET_SIZE: usize = 512;
/// Size for control transfers on endpoint 0.
const EP0_CONTROL_PACKET_SIZE: usize = 64;
/// The USB GPT timer we use to (infrequently) check for data.
const GPT_INSTANCE: imxrt_usbd::gpt::Instance = imxrt_usbd::gpt::Instance::Gpt0;

struct Backend<'a> {
    class: Class<'a>,
    device: Device<'a>,
    configured: bool,
    stdout_queue: &'static BlockQueue,
}

impl Backend<'_> {
    fn poll(&mut self) {
        // Is there a CDC class event, like a completed transfer? If so, check
        // the consumer immediately, even if a timer hasn't expired.
        //
        // Checking the consumer on class traffic lets the driver burst out data.
        // Suppose the user wants to use the USB GPT timer, and they configure a very
        // long interval. That interval expires, and we see tons of data in the consumer.
        // We should write that out as fast as possible, even if the timer hasn't elapsed.
        // That's the behavior provided by the class_event flag.
        let class_event = self.device.poll(&mut [&mut self.class]);
        let timer_event = self.device.bus().gpt_mut(GPT_INSTANCE, |gpt| {
            let mut elapsed = false;
            while gpt.is_elapsed() {
                gpt.clear_elapsed();
                elapsed = true;
            }
            // Simulate a timer event if the timer is not running.
            //
            // If the timer is not running, its because the user disabled interrupts,
            // and they're using their own timer / polling loop. There might not always
            // be a class traffic (transfer complete) event when the user polls, so
            // signaling true allows the poll to check the consumer for new data and
            // send it.
            //
            // If the timer is running, checking the consumer depends on the elapsed
            // timer.
            elapsed || !gpt.is_running()
        });
        let check_consumer = class_event || timer_event;

        if self.device.state() != UsbDeviceState::Configured {
            if self.configured {
                // Turn off the timer, but only if we were previously configured.
                self.device.bus().gpt_mut(GPT_INSTANCE, |gpt| gpt.stop());
            }
            self.configured = false;
            // We can't use the class if we're not configured,
            // so bail out here.
            return;
        }

        // We're now configured. Are we newly configured?
        if !self.configured {
            // Must call this when we transition into configured.
            self.device.bus().configure();
            self.device.bus().gpt_mut(GPT_INSTANCE, |gpt| {
                // There's no need for a timer if interrupts are disabled.
                // If the user disabled USB interrupts and decided to poll this
                // from another timer, this USB timer could unnecessarily block
                // that timer from checking the consumer queue.
                if gpt.is_interrupt_enabled() {
                    gpt.run()
                }
            });
            self.configured = true;
        }

        self.class.read_packet(&mut []).ok();

        if check_consumer && let Some(block) = self.stdout_queue.try_receive() {
            // Ignoring errors here... and we're not really utilizing the full
            // width of the endpoint buffer. Let's see how long we can get
            // away with this...
            let _ = self.class.write_packet(&block.byte_buffer[..block.len]);
        }
    }
}

fn run<const N: u8>(
    peripherals: imxrt_usbd::Instances<N>,
    interrupt: imxrt_ral::Interrupt,
    stdout_queue: &'static BlockQueue,
) -> ! {
    let events = Semaphore::create(Pin::static_ref(&EVENTS), &Default::default()).unwrap();

    let bus = {
        // Safety: we ensure that the bus, class, and all other related USB objects
        // are accessed in poll(). poll() is not reentrant, so there's no racing
        // occuring across executing contexts.
        let bus = unsafe {
            imxrt_usbd::BusAdapter::without_critical_sections(
                peripherals,
                &ENDPOINT_MEMORY,
                &ENDPOINT_STATE,
                imxrt_usbd::Speed::High,
            )
        };
        bus.set_interrupts(true);
        bus.gpt_mut(GPT_INSTANCE, |gpt| {
            gpt.stop();
            gpt.clear_elapsed();
            gpt.set_interrupt_enabled(true);
            gpt.set_mode(imxrt_usbd::gpt::Mode::Repeat);
            gpt.set_load(4_000);
            gpt.reset();
        });
        usb_device::bus::UsbBusAllocator::new(bus)
    };
    let class = usbd_serial::CdcAcmClass::new(&bus, MAX_PACKET_SIZE as u16);

    let device = usb_device::device::UsbDeviceBuilder::new(&bus, VID_PID)
        .strings(&[usb_device::device::StringDescriptors::default().product(PRODUCT)])
        .unwrap()
        .device_class(usbd_serial::USB_CLASS_CDC)
        .max_packet_size_0(EP0_CONTROL_PACKET_SIZE as u8)
        .unwrap()
        .build();

    // Not sure which endpoints the CDC ACM class will pick,
    // so enable the setting for all non-zero endpoints.
    for idx in 1..8 {
        for dir in &[usb_device::UsbDirection::In, usb_device::UsbDirection::Out] {
            let ep_addr = usb_device::endpoint::EndpointAddress::from_parts(idx, *dir);
            // CDC class requires that we send the ZLP.
            // Let the hardware do that for us.
            device.bus().enable_zlt(ep_addr);
        }
    }

    let mut backend = Backend {
        class,
        device,
        configured: false,
        stdout_queue,
    };

    loop {
        backend.poll();
        unsafe { cortex_m::peripheral::NVIC::unmask(interrupt) };
        events.get().unwrap();
    }
}

use core::pin::Pin;
use fusible::semaphore::{Semaphore, SemaphoreContext};

static EVENTS: SemaphoreContext = Semaphore::context();

#[inline]
pub fn on_interrupt(intr: imxrt_ral::Interrupt) {
    cortex_m::peripheral::NVIC::mask(intr);
    if let Some(events) = Pin::static_ref(&EVENTS).try_created() {
        events.put();
    }
}

use fusible::thread::{StaticStack, Thread, ThreadContext};

static STACK: StaticStack<2048> = StaticStack::new();
static THREAD: ThreadContext = Thread::context();

pub fn start<const N: u8>(peripherals: imxrt_usbd::Instances<N>, interrupt: imxrt_ral::Interrupt) {
    BlockPool::create(
        Pin::static_ref(&BLOCK_POOL),
        BLOCK_STORAGE.take().unwrap(),
        &Default::default(),
    )
    .unwrap();

    let stdout_queue = Queue::create(
        Pin::static_ref(&STDOUT_QUEUE),
        STDOUT_STORAGE.take().unwrap(),
        &Default::default(),
    )
    .unwrap();

    Thread::create(
        Pin::static_ref(&THREAD),
        STACK.take().unwrap(),
        &Default::default(),
        move || run(peripherals, interrupt, stdout_queue),
    )
    .unwrap();
}

use fusible::block_pool::{Block, BlockPool, BlockPoolContext, StaticBlocks};

type ByteBuffer = [u8; MAX_PACKET_SIZE];

struct BlockWithLen {
    len: usize,
    byte_buffer: ByteBuffer,
}

impl BlockWithLen {
    fn copy_of(buf: &[u8]) -> Self {
        let mut byte_buffer: ByteBuffer = [0_u8; _];
        let len = byte_buffer.len().min(buf.len());
        byte_buffer[..len].copy_from_slice(&buf[..len]);
        Self { len, byte_buffer }
    }
}

const BLOCK_COUNT: usize = 4;
static BLOCK_STORAGE: StaticBlocks<BlockWithLen, BLOCK_COUNT> = StaticBlocks::new();
static BLOCK_POOL: BlockPoolContext<BlockWithLen> = BlockPool::context();

type ByteBufferBlock = Block<'static, BlockWithLen>;

use fusible::queue::{Queue, QueueContext, StaticQueueSlots};

static STDOUT_STORAGE: StaticQueueSlots<ByteBufferBlock, BLOCK_COUNT> = StaticQueueSlots::new();
static STDOUT_QUEUE: QueueContext<ByteBufferBlock> = Queue::context();

type BlockQueue = Queue<ByteBufferBlock>;

fn stdout_write(buf: &[u8]) -> usize {
    let mut written = 0;

    let Some(block_pool) = Pin::static_ref(&BLOCK_POOL).try_created() else {
        return written;
    };

    let Some(stdout_queue) = Pin::static_ref(&STDOUT_QUEUE).try_created() else {
        return written;
    };

    for chunk in buf.chunks(MAX_PACKET_SIZE) {
        let Some(block) = block_pool.try_allocate(|| {
            let block = BlockWithLen::copy_of(&chunk);
            written += block.len;
            block
        }) else {
            break;
        };

        if stdout_queue.try_send(block).is_some() {
            break;
        }
    }

    written
}

#[unsafe(no_mangle)]
fn __rust_threadx_stdout_write(buf: &[u8]) -> usize {
    stdout_write(buf)
}

#[unsafe(no_mangle)]
fn __rust_threadx_stderr_write(buf: &[u8]) -> usize {
    stdout_write(buf)
}

#[unsafe(no_mangle)]
fn __rust_threadx_stdin_read(_: &mut [u8]) -> usize {
    0 // Haven't gotten this working yet.
}