use core::{
    future::poll_fn,
    pin::Pin,
    task::{Context, Poll, Waker},
    time::Duration,
};

use alloc::vec::Vec;
use futures_util::{Future, FutureExt};
use kernel_arch::{Architecture, ArchitectureImpl};
use libk_util::sync::IrqSafeSpinlock;
use yggdrasil_abi::{error::Error, time::SystemTime};

use crate::time::monotonic_time;

// 2..32ms, tick every 2ms
static SHORT_TERM_SLEEPS: IrqSafeSpinlock<TimerWheel<16>> =
    IrqSafeSpinlock::new(TimerWheel::new(0, 2));
// 32ms..288ms, tick every 8ms
static LONG_TERM_SLEEPS: IrqSafeSpinlock<TimerWheel<32>> =
    IrqSafeSpinlock::new(TimerWheel::new(4, 8));

// A simple timer wheel inspired by
// https://www.snellman.net/blog/archive/2016-07-27-ratas-hierarchical-timer-wheel/
// The wheel consists of a ring with slots. When some task wants to sleep for N ticks,
// it adds itself to a list N steps after current tick (modulo STEPS). The whell also contains
// a sort of interval to "scale" ticks for longer-term sleeps.
struct TimerWheel<const STEPS: usize> {
    // TODO maybe use something better than Vec
    ring: [Vec<Waker>; STEPS],
    tick: u64,
    real_tick: u64,
    base: u64,
    interval: u64,
}

impl<const STEPS: usize> TimerWheel<STEPS> {
    pub const fn new(base: u64, interval: u64) -> Self {
        Self {
            ring: [const { Vec::new() }; STEPS],
            tick: 0,
            real_tick: 0,
            base,
            interval,
        }
    }

    pub fn tick(&mut self) {
        self.real_tick += 1;
        if self.real_tick == self.interval {
            self.tick = self.tick.wrapping_add(1);
            for waker in self.ring[(self.tick % (STEPS as u64)) as usize].drain(..) {
                waker.wake();
            }
            self.real_tick = 0;
        }
    }

    pub fn wake_after(&mut self, ticks: u64, waker: &Waker) {
        debug_assert!(
            (ticks / self.interval) - self.base <= STEPS as u64,
            "Invalid tick count: ticks={ticks}, interval={}, base={}, STEPS={STEPS}",
            self.interval,
            self.base
        );
        let ticks = ticks.max(1);
        self.ring[((self.tick + ticks / self.interval - self.base) % (STEPS as u64)) as usize]
            .push(waker.clone());
    }
}

fn register_timeout(duration: Duration, waker: &Waker) {
    let nticks = duration.as_millis().clamp(2, 288) as u64;

    if nticks < 32 {
        SHORT_TERM_SLEEPS.lock().wake_after(nticks, waker);
    } else {
        LONG_TERM_SLEEPS.lock().wake_after(nticks, waker);
    }
}

pub struct SleepFuture {
    deadline: SystemTime,
}

impl SleepFuture {
    pub fn remaining(&self) -> Option<Duration> {
        let now = monotonic_time();
        match self.deadline.checked_sub_time(&now) {
            Some(duration) if !duration.is_zero() => Some(duration),
            _ => None,
        }
    }
}

impl Future for SleepFuture {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        match self.remaining() {
            // Pending
            Some(duration) => {
                register_timeout(duration, cx.waker());
                Poll::Pending
            }
            // Ran out
            _ => Poll::Ready(()),
        }
    }
}

/// Suspends the task until given duration passes.
pub fn sleep(duration: Duration) -> SleepFuture {
    let now = monotonic_time();
    let deadline = now + duration;

    SleepFuture { deadline }
}

/// Suspends the task until the given deadline comes.
pub fn sleep_until(deadline: SystemTime) -> SleepFuture {
    SleepFuture { deadline }
}

/// Wait for at least a certain duration has passed.
///
/// # Note
///
/// This function unmasks IRQs. Do not use it where interrupts can break something.
pub fn psleep(duration: Duration) {
    // Unlock IRQs here, otherwise no timer events will be received and monotonic_time() will stay
    // the same
    let now = monotonic_time();
    let deadline = now + duration;

    let old_irq_mask = ArchitectureImpl::interrupt_mask();
    unsafe { ArchitectureImpl::set_interrupt_mask(false) };

    while monotonic_time() < deadline {
        core::hint::spin_loop();
    }

    // Restore IRQ mask
    unsafe { ArchitectureImpl::set_interrupt_mask(old_irq_mask) };
}

/// Wait for some condition to become true, or for the timeout to expire.
///
/// # Returns
///
/// * [Error::TimedOut] if the condition didn't become `true` within a given deadline.
///
/// # Note
///
/// This function unmasks IRQs. Do not use it where interrupts can break something.
pub fn pwait<P: Fn() -> bool>(
    timeout: Duration,
    poll_rate: Duration,
    predicate: P,
) -> Result<(), Error> {
    let now = monotonic_time();
    let deadline = now + timeout;

    loop {
        if predicate() {
            return Ok(());
        }
        if monotonic_time() > deadline {
            return Err(Error::TimedOut);
        }
        psleep(poll_rate);
    }
}

/// Wait for some condition to become true, fail, or for the timeout to expire.
///
/// # Returns
///
/// * [Error] if the predicate returns an error.
/// * [Error::TimedOut] if the condition didn't become `true` within a given deadline.
///
/// # Note
///
/// This function unmasks IRQs. Do not use it where interrupts can break something.
pub fn pwait_try<P: Fn() -> Result<bool, Error>>(
    timeout: Duration,
    poll_rate: Duration,
    predicate: P,
) -> Result<(), Error> {
    let now = monotonic_time();
    let deadline = now + timeout;

    loop {
        if predicate()? {
            return Ok(());
        }
        if monotonic_time() > deadline {
            return Err(Error::TimedOut);
        }
        psleep(poll_rate);
    }
}

pub fn with_timeout<'a, T: 'a, F: Future<Output = T> + Send + 'a>(
    fut: F,
    timeout: Duration,
) -> impl Future<Output = Result<T, Error>> + Send + 'a {
    let mut fut = fut.boxed();
    let mut sleep_fut = sleep(timeout).boxed();

    poll_fn(move |cx| {
        let (timeout, output) = (sleep_fut.poll_unpin(cx), fut.poll_unpin(cx));
        if let Poll::Ready(output) = output {
            Poll::Ready(Ok(output))
        } else if timeout.is_ready() {
            Poll::Ready(Err(Error::TimedOut))
        } else {
            Poll::Pending
        }
    })
}

pub async fn maybe_timeout<T, F: Future<Output = T> + Send>(
    fut: F,
    timeout: Option<Duration>,
) -> Result<T, Error> {
    match timeout {
        Some(timeout) => with_timeout(fut, timeout).await,
        None => Ok(fut.await),
    }
}

/// Updates the runtime's time
pub fn tick() {
    SHORT_TERM_SLEEPS.lock().tick();
    LONG_TERM_SLEEPS.lock().tick();
}
