#![cfg_attr(target_os = "yggdrasil", feature(yggdrasil_os, rustc_private))]
use std::{
    cell::OnceCell,
    collections::HashMap,
    future::Future,
    os::fd::{AsRawFd, RawFd},
    pin::Pin,
    sync::{Arc, Mutex},
    task::{Context, Poll, Waker},
    thread::JoinHandle,
    time::Duration,
};

use cross::io::{Poll as PollChannel, TimerFd};
use futures_util::task::{waker_ref, ArcWake};
use spmc::{Receiver, Sender};

pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;

pub struct Executor {
    _id: usize,
    queue: Arc<Receiver<Arc<Task>>>,
}

pub struct Task {
    future: Mutex<Option<BoxFuture<'static, ()>>>,
    runtime: Arc<RuntimeInner>,
}

struct ReactorInner {
    poll: PollChannel,
    event_map: HashMap<RawFd, Waker>,
}

pub struct Reactor {
    // TODO this lock will become very contended
    inner: Mutex<ReactorInner>,
}

struct RuntimeInner {
    // TODO this lock will become very contended
    sender: Mutex<Sender<Arc<Task>>>,
    reactor: Arc<Reactor>
}

pub struct Runtime {
    inner: Arc<RuntimeInner>,
    executors: Vec<JoinHandle<()>>,
    reactor: JoinHandle<()>,
}

// TODO use global timer instead of creating a timer for each sleep
pub struct Sleep {
    timer: TimerFd,
}

impl Future for Sleep {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let expired = self.timer.is_expired().unwrap();
        if expired {
            return Poll::Ready(());
        }
        let reactor = &RuntimeInner::get().reactor;
        reactor.register(&self.timer, cx.waker().clone());
        Poll::Pending
    }
}

impl Executor {
    pub fn new(id: usize, queue: Arc<Receiver<Arc<Task>>>) -> Self {
        Self { _id: id, queue }
    }

    pub fn run(self) {
        loop {
            let task = self.queue.recv().unwrap();
            let mut slot = task.future.lock().unwrap();

            if let Some(mut future) = slot.take() {
                let waker = waker_ref(&task);
                let context = &mut Context::from_waker(&waker);

                if future.as_mut().poll(context).is_pending() {
                    *slot = Some(future);
                }
            }
        }
    }
}

impl ArcWake for Task {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        arc_self
            .runtime
            .sender
            .lock()
            .unwrap()
            .send(arc_self.clone())
            .unwrap();
    }
}

thread_local! {
    static RUNTIME: OnceCell<Arc<RuntimeInner>> = const { OnceCell::new() };
}

impl RuntimeInner {
    fn try_get() -> Option<Arc<RuntimeInner>> {
        RUNTIME.with(|r| r.get().cloned())
    }

    fn get() -> Arc<RuntimeInner> {
        Self::try_get().expect("Runtime::get() called from a thread with no runtime")
    }
}

impl Reactor {
    fn new() -> Self {
        Self {
            inner: Mutex::new(ReactorInner {
                poll: PollChannel::new().unwrap(),
                event_map: HashMap::new(),
            }),
        }
    }

    fn run(self: Arc<Self>) {
        loop {
            // TODO this is a hack to prevent the lock from becoming fully contended by Reactor
            std::thread::sleep(Duration::from_millis(100));

            let mut inner = self.inner.lock().unwrap();
            let event = inner.poll.wait(Some(Duration::from_millis(100))).unwrap();
            if let Some(fd) = event {
                inner.poll.remove(&fd).ok();
                if let Some(task) = inner.event_map.remove(&fd) {
                    task.wake();
                }
            }
        }
    }

    fn register<F: AsRawFd>(&self, interest: &F, waker: Waker) {
        let mut inner = self.inner.lock().unwrap();
        let fd = interest.as_raw_fd();
        // TODO have multiple tasks be able to poll the same fd?
        assert!(inner.event_map.insert(fd, waker).is_none());
        inner.poll.add(interest).unwrap();
    }
}

impl Runtime {
    pub fn new(threads: usize) -> Self {
        let mut executors = vec![];
        let (sender, receiver) = spmc::channel();
        let reactor = Arc::new(Reactor::new());
        let inner = Arc::new(RuntimeInner {
            reactor: reactor.clone(),
            sender: Mutex::new(sender),
        });
        let receiver = Arc::new(receiver);

        // TODO add an option to run reactor on main thread?
        let reactor = std::thread::spawn(move || reactor.run());

        for id in 0..threads {
            let receiver = receiver.clone();
            let inner = inner.clone();
            executors.push(std::thread::spawn(move || {
                // Initialize local runtime ref
                RUNTIME.with(|r| {
                    if r.set(inner).is_err() {
                        panic!("Couldn't initialize thread runtime");
                    }
                });

                Executor::new(id, receiver).run();
            }));
        }

        Self {
            inner,
            executors,
            reactor,
        }
    }

    pub fn spawn<F: Future<Output = ()> + Send + 'static>(&self, future: F) {
        spawn_into(self.inner.clone(), future)
    }

    pub fn join(self) {
        for executor in self.executors {
            executor.join().ok();
        }
        self.reactor.join().ok();
    }
}

fn spawn_into<F: Future<Output = ()> + Send + 'static>(runtime: Arc<RuntimeInner>, future: F) {
    let task = Arc::new(Task {
        future: Mutex::new(Some(Box::pin(future))),
        runtime: runtime.clone(),
    });
    runtime.sender.lock().unwrap().send(task).unwrap();
}

pub fn spawn<F: Future<Output = ()> + Send + 'static>(future: F) {
    let runtime = RuntimeInner::get();
    spawn_into(runtime, future)
}

pub fn sleep(timeout: Duration) -> Sleep {
    let mut timer = TimerFd::new().unwrap();
    timer.start(timeout).unwrap();
    Sleep { timer }
}
