use std::{
    pin::Pin,
    sync::{
        atomic::{AtomicU64, AtomicUsize, Ordering},
        Arc,
    },
};

use futures::Future;
use tokio::sync::Notify;

use crate::{
    builder::Config,
    chan::{BusSender, Receiver, Sender},
    handler::HandlerSpawner,
    message::Msg,
    Builder, BusInner, Error, Handler, Message,
};

#[derive(Default)]
pub(crate) struct TaskCounter {
    pub(crate) running: AtomicUsize,
    notify: Notify,
}

pub(crate) struct TaskCounterLease<S: Fn() -> bool> {
    need_notify: S,
    counter: Arc<TaskCounter>,
}

impl<S: Fn() -> bool> Drop for TaskCounterLease<S> {
    fn drop(&mut self) {
        let notify = (self.need_notify)();
        let prev = self.counter.running.fetch_sub(1, Ordering::Relaxed);

        if notify && prev == 1 {
            self.counter.notify.notify_waiters();
        }
    }
}

impl<S: Fn() -> bool> TaskCounterLease<S> {
    fn new(counter: Arc<TaskCounter>, need_notify: S) -> Self {
        counter.running.fetch_add(1, Ordering::Relaxed);

        Self {
            counter,
            need_notify,
        }
    }
}

impl TaskCounter {
    pub fn lease_unit<S: Fn() -> bool>(self: Arc<Self>, need_notify: S) -> TaskCounterLease<S> {
        TaskCounterLease::new(self, need_notify)
    }

    #[inline]
    pub async fn wait(&self) {
        self.notify.notified().await
    }
}

pub(crate) trait TaskSpawner<M: Message>: Send + Sync {
    fn config(&self, stream_id: u32) -> Config;
    fn is_producer(&self) -> bool;

    #[allow(clippy::too_many_arguments)]
    fn spawn_task(
        &self,
        rx: Receiver<Msg<M>>,
        stream_id: u32,
        task_id: u32,
        abort: Arc<Notify>,
        task_counter: Arc<TaskCounter>,
        spawn_counter: Arc<TaskCounter>,
        index_counter: Arc<AtomicU64>,
        bus: Arc<BusInner>,
    ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>>;
}

pub(crate) struct TaskSpawnerWrapper<M: Message> {
    inner: Arc<dyn TaskSpawner<M>>,
}

impl<M: Message> Clone for TaskSpawnerWrapper<M> {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
        }
    }
}

impl<M: Message> TaskSpawnerWrapper<M> {
    pub fn from_handler<B: Builder<M> + 'static>(builder: B) -> Self
    where
        B::Context: Handler<M>,
    {
        Self {
            inner: Arc::new(HandlerSpawner::new(builder)) as _,
        }
    }

    #[inline]
    #[allow(clippy::too_many_arguments)]
    pub async fn spawn_task(
        &self,
        (tx, rx): (Sender<Msg<M>>, Receiver<Msg<M>>),
        stream_id: u32,
        task_id: u32,
        abort: Arc<Notify>,
        task_counter: Arc<TaskCounter>,
        spawn_counter: Arc<TaskCounter>,
        index_counter: Arc<AtomicU64>,
        bus: Arc<BusInner>,
    ) -> Result<BusSender<M>, Error> {
        self.inner
            .spawn_task(
                rx,
                stream_id,
                task_id,
                abort,
                task_counter,
                spawn_counter,
                index_counter,
                bus,
            )
            .await?;

        Ok(BusSender::new(self.inner.is_producer(), tx))
    }

    #[inline]
    pub(crate) fn config(&self, stream_id: u32) -> Config {
        self.inner.config(stream_id)
    }
}