use core::task::{Context, Poll, Waker};

use alloc::{boxed::Box, format, sync::Arc};
use futures_util::{task::waker_ref, Future};
use kernel_arch::task::TaskContext;
use libk_util::waker::WakeWeak;
use yggdrasil_abi::error::Error;

use crate::task::{thread::Thread, TaskContextImpl};

use super::{
    task::{Task, Termination},
    task_queue,
};

/// Pushes a task into the executor's queue
pub fn enqueue(task: Arc<Task>) -> Result<(), Error> {
    task_queue::push_task(task)
}

/// Spawns a background worker to execute the tasks from the global queue
pub fn spawn_async_worker(index: usize) -> Result<(), Error> {
    let name = format!("[async-worker-{index}]");

    let thread = Thread::new_kthread(
        name,
        TaskContextImpl::kernel_closure(move || loop {
            let task = task_queue::pop_task().unwrap();
            let mut future_slot = task.future.lock();

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

                if future.as_mut().poll(context).is_pending() {
                    *future_slot = Some(future);
                }
            }
        })?,
    );

    thread.enqueue();

    Ok(())
}

/// Creates a new task for the [Future] and queues it for execution in background
pub fn spawn<T: Termination, F: Future<Output = T> + Send + 'static>(
    future: F,
) -> Result<(), Error> {
    enqueue(Task::new(future))
}

/// Runs a [Future] to its completion on the current thread
pub fn run_to_completion<'a, T, F: Future<Output = T> + 'a>(future: F) -> Result<T, Error> {
    let mut future = Box::pin(future);

    // Make a weak ref for the waker
    let (entered, weak) = {
        let current = Thread::current();
        (current.id, current.downgrade())
    };

    let waker = WakeWeak::weak_waker(weak.clone());
    let waker = unsafe { Waker::from_raw(waker) };

    let result = loop {
        let context = &mut Context::from_waker(&waker);

        match future.as_mut().poll(context) {
            Poll::Ready(value) => break Ok(value),
            Poll::Pending => {
                if let Some(thread) = unsafe { Thread::upgrade(&weak) } {
                    if let Err(error) = thread.suspend() {
                        log::warn!("Future cancelled: {error:?}");
                        break Err(error);
                    }
                }
            }
        }
    };

    drop(future);

    let current = Thread::current().id;
    assert_eq!(current, entered);

    result
}
