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

use core::future::Future;
use core::mem;
use core::pin::Pin;
use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};

use cortex_m_semihosting::{debug, hprintln};
use panic_semihosting as _;
use systick_monotonic::*;

// NOTES:
//
// - Async tasks cannot have `#[lock_free]` resources, as they can interleve and each async
//   task can have a mutable reference stored.
// - Spawning an async task equates to it being polled at least once.
// - ...

#[rtic::app(device = lm3s6965, dispatchers = [SSI0], peripherals = true)]
mod app {
    use crate::*;

    pub type AppInstant = <Systick<100> as rtic::Monotonic>::Instant;
    pub type AppDuration = <Systick<100> as rtic::Monotonic>::Duration;

    #[shared]
    struct Shared {
        s: u32,
    }

    #[local]
    struct Local {}

    #[monotonic(binds = SysTick, default = true)]
    type MyMono = Systick<100>;

    #[init]
    fn init(cx: init::Context) -> (Shared, Local, init::Monotonics) {
        hprintln!("init").unwrap();
        task_executor::spawn().unwrap();

        (
            Shared { s: 0 },
            Local {},
            init::Monotonics(Systick::new(cx.core.SYST, 12_000_000)),
        )
    }

    #[idle]
    fn idle(_: idle::Context) -> ! {
        // debug::exit(debug::EXIT_SUCCESS);
        loop {
            // hprintln!("idle");
            cortex_m::asm::wfi(); // put the MCU in sleep mode until interrupt occurs
        }
    }

    // TODO: This should be the task, that is understood by the `syntax` proc-macro
    // #[task(priority = 2)]
    async fn task(cx: task_executor::Context<'_>) {
        #[allow(unused_imports)]
        use rtic::mutex_prelude::*;

        hprintln!("delay long time").ok();

        let fut = Delay::spawn(2500.millis());

        hprintln!("we have just created the future").ok();
        fut.await;
        hprintln!("long delay done").ok();

        hprintln!("delay short time").ok();
        sleep(500.millis()).await;
        hprintln!("short delay done").ok();

        hprintln!("test timeout").ok();
        let res = timeout(NeverEndingFuture {}, 1.secs()).await;
        hprintln!("timeout done: {:?}", res).ok();

        hprintln!("test timeout 2").ok();
        let res = timeout(Delay::spawn(500.millis()), 1.secs()).await;
        hprintln!("timeout done 2: {:?}", res).ok();

        debug::exit(debug::EXIT_SUCCESS);
    }

    //////////////////////////////////////////////
    // BEGIN BOILERPLATE
    //////////////////////////////////////////////
    type F = impl Future + 'static;
    static mut TASK: AsyncTaskExecutor<F> = AsyncTaskExecutor::new();

    // TODO: This should be a special case codegen for the `dispatcher`, which runs
    //       in the dispatcher. Not as its own task, this is just to make it work
    //       in this example.
    #[task(shared = [s])]
    fn task_executor(cx: task_executor::Context) {
        let task_storage = unsafe { &mut TASK };
        match task_storage {
            AsyncTaskExecutor::Idle => {
                // TODO: The context generated for async tasks need 'static lifetime,
                // use `mem::transmute` for now until codegen is fixed
                //
                // TODO: Check if there is some way to not need 'static lifetime
                hprintln!("    task_executor spawn").ok();
                task_storage.spawn(|| task(unsafe { mem::transmute(cx) }));
                task_executor::spawn().ok();
            }
            _ => {
                hprintln!("    task_executor run").ok();
                task_storage.poll(|| {
                    task_executor::spawn().ok();
                });
            }
        };
    }
    //////////////////////////////////////////////
    // END BOILERPLATE
    //////////////////////////////////////////////

    // TODO: This is generated by the `delay` impl, it needs a capacity equal or grater
    //       than the number of async tasks in the system. Should more likely be a part
    //       of the monotonic codegen, not its own task.
    #[task(capacity = 12)]
    fn delay_handler(_: delay_handler::Context, waker: Waker) {
        waker.wake();
    }
}

//=============
// Waker

static WAKER_VTABLE: RawWakerVTable =
    RawWakerVTable::new(waker_clone, waker_wake, waker_wake, waker_drop);

unsafe fn waker_clone(p: *const ()) -> RawWaker {
    RawWaker::new(p, &WAKER_VTABLE)
}

unsafe fn waker_wake(p: *const ()) {
    // The only thing we need from a waker is the function to call to pend the async
    // dispatcher.
    let f: fn() = mem::transmute(p);
    f();
}

unsafe fn waker_drop(_: *const ()) {
    // nop
}

//============
// AsyncTaskExecutor

enum AsyncTaskExecutor<F: Future + 'static> {
    Idle,
    Running(F),
}

impl<F: Future + 'static> AsyncTaskExecutor<F> {
    const fn new() -> Self {
        Self::Idle
    }

    fn spawn(&mut self, future: impl FnOnce() -> F) {
        *self = AsyncTaskExecutor::Running(future());
    }

    fn poll(&mut self, wake: fn()) {
        match self {
            AsyncTaskExecutor::Idle => {}
            AsyncTaskExecutor::Running(future) => unsafe {
                let waker_data: *const () = mem::transmute(wake);
                let waker = Waker::from_raw(RawWaker::new(waker_data, &WAKER_VTABLE));
                let mut cx = Context::from_waker(&waker);
                let future = Pin::new_unchecked(future);

                match future.poll(&mut cx) {
                    Poll::Ready(_) => {
                        *self = AsyncTaskExecutor::Idle;
                        hprintln!("    task_executor idle").ok();
                    }
                    Poll::Pending => {}
                };
            },
        }
    }
}

//=============
// Delay

pub struct Delay {
    until: crate::app::AppInstant,
}

impl Delay {
    pub fn spawn(duration: crate::app::AppDuration) -> Self {
        let until = crate::app::monotonics::now() + duration;

        Delay { until }
    }
}

#[inline(always)]
pub fn sleep(duration: crate::app::AppDuration) -> Delay {
    Delay::spawn(duration)
}

impl Future for Delay {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let s = self.as_mut();
        let now = crate::app::monotonics::now();

        hprintln!("    poll Delay").ok();

        if now >= s.until {
            Poll::Ready(())
        } else {
            let waker = cx.waker().clone();
            crate::app::delay_handler::spawn_after(s.until - now, waker).ok();

            Poll::Pending
        }
    }
}

//=============
// Timeout future

#[derive(Copy, Clone, Debug)]
pub struct TimeoutError;

pub struct Timeout<F: Future> {
    future: F,
    until: crate::app::AppInstant,
    cancel_handle: Option<crate::app::delay_handler::SpawnHandle>,
}

impl<F> Timeout<F>
where
    F: Future,
{
    pub fn timeout(future: F, duration: crate::app::AppDuration) -> Self {
        let until = crate::app::monotonics::now() + duration;
        Self {
            future,
            until,
            cancel_handle: None,
        }
    }
}

#[inline(always)]
pub fn timeout<F: Future>(future: F, duration: crate::app::AppDuration) -> Timeout<F> {
    Timeout::timeout(future, duration)
}

impl<F> Future for Timeout<F>
where
    F: Future,
{
    type Output = Result<F::Output, TimeoutError>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let now = crate::app::monotonics::now();

        // SAFETY: We don't move the underlying pinned value.
        let mut s = unsafe { self.get_unchecked_mut() };
        let future = unsafe { Pin::new_unchecked(&mut s.future) };

        hprintln!("    poll Timeout").ok();

        match future.poll(cx) {
            Poll::Ready(r) => {
                if let Some(ch) = s.cancel_handle.take() {
                    ch.cancel().ok();
                }

                Poll::Ready(Ok(r))
            }
            Poll::Pending => {
                if now >= s.until {
                    Poll::Ready(Err(TimeoutError))
                } else if s.cancel_handle.is_none() {
                    let waker = cx.waker().clone();
                    let sh = crate::app::delay_handler::spawn_after(s.until - now, waker)
                        .expect("Internal RTIC bug, this should never fail");
                    s.cancel_handle = Some(sh);

                    Poll::Pending
                } else {
                    Poll::Pending
                }
            }
        }
    }
}

pub struct NeverEndingFuture {}

impl Future for NeverEndingFuture {
    type Output = ();

    fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
        // Never finish
        hprintln!("    polling NeverEndingFuture").ok();
        Poll::Pending
    }
}

//=============
// Async SPI driver

// #[task]
async fn test_spi(async_spi_driver: &mut AsyncSpi) {
    let transfer = Transaction {
        buf: [0; 16],
        n_write: 1,
        n_read: 5,
    };

    let ret = async_spi_driver.transfer(transfer).await;

    // do_something(ret);
}

/// A DMA transaction.
///
/// NOTE: Don't leak this `Future`, if you do there is immediate UB!
struct Transaction {
    pub buf: [u8; 16],
    pub n_write: usize,
    pub n_read: usize,
}

struct AsyncSpi {
    transaction: Option<Transaction>,
    queue: heapless::spsc::Queue<Waker, 8>,
}

impl AsyncSpi {
    pub fn transfer(&mut self, transfer: Transaction) -> AsyncSpiTransaction {
        todo!()
    }
}

struct AsyncSpiTransaction {
    // ...
}

impl Future for AsyncSpiTransaction {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        todo!()
    }
}