#![doc = include_str!("../README.md")]
#![deny(missing_docs)]

mod drop_notifier;
mod executor;
mod notify;
mod queue;
mod selector;
mod spsc;

use std::{
    future::Future,
    num::{NonZeroU16, NonZeroUsize},
    sync::{atomic::AtomicU64, Arc, Mutex},
    thread::JoinHandle,
    time::Instant,
};

use drop_notifier::{DropListener, DropNotifier};
use executor::block_on;
use queue::Queue;
use selector::select;

/// Configuration builder for the CpuPool
#[derive(Debug)]
pub struct Config {
    name: &'static str,
    buffer_multiplier: usize,
    min_threads: u16,
    max_threads: u16,
}

impl Config {
    /// Create a new configuration builder with the default configuration
    pub fn new() -> Self {
        Config {
            name: "cpupool",
            buffer_multiplier: 8,
            min_threads: 1,
            max_threads: 4,
        }
    }

    /// Set the name for the CpuPool
    ///
    /// This is used for setting the names of spawned threads
    ///
    /// default: `"cpupool"`
    ///
    /// Example:
    /// ```rust
    /// # use async_cpupool::Config;
    /// Config::new().name("sig-pool");
    /// ```
    pub fn name(mut self, name: &'static str) -> Self {
        self.name = name;
        self
    }

    /// Set the multiplier for the internal queue's buffer size
    ///
    /// This value must be at least 1. the buffer's size will be equal to `max_threads * buffer_multiplier`
    ///
    /// default: `8`
    ///
    /// Example:
    /// ```rust
    /// # use async_cpupool::Config;
    /// Config::new().buffer_multiplier(4);
    /// ```
    pub fn buffer_multiplier(mut self, buffer_multiplier: usize) -> Self {
        self.buffer_multiplier = buffer_multiplier;
        self
    }

    /// Set the minimum allowed number of running threads
    ///
    /// When there is little work to do, threads will be reaped until just this number remain
    ///
    /// default: `1`
    ///
    /// Example:
    /// ```rust
    /// # use async_cpupool::Config;
    /// Config::new().min_threads(2);
    /// ```
    pub fn min_threads(mut self, min_threads: u16) -> Self {
        self.min_threads = min_threads;
        self
    }

    /// Set the maximum allowed number of running threads
    ///
    /// When the threadpool is under load, threads will be spawned until this limit is reached
    ///
    /// default: `4`
    ///
    /// Example:
    /// ```rust
    /// # use async_cpupool::Config;
    /// Config::new().max_threads(16);
    /// ```
    pub fn max_threads(mut self, max_threads: u16) -> Self {
        self.max_threads = max_threads;
        self
    }

    /// Create a CpuPool with the given configuration, spawning `min_threads` threads
    ///
    /// This will error if `min_threads` is greater than `max_threads`, or if `buffer_multiplier`,
    /// `max_threads`, or `min_threads` are `0`
    ///
    /// Example:
    /// ```rust
    /// # use async_cpupool::Config;
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let pool = Config::new()
    ///     .name("sig-pool")
    ///     .min_threads(4)
    ///     .max_threads(16)
    ///     .buffer_multiplier(2)
    ///     .build()?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn build(self) -> Result<CpuPool, ConfigError> {
        let Config {
            name,
            buffer_multiplier,
            min_threads,
            max_threads,
        } = self;

        if max_threads < min_threads {
            return Err(ConfigError::ThreadCount);
        }

        let buffer_multiplier = buffer_multiplier
            .try_into()
            .map_err(|_| ConfigError::BufferMultiplier)?;

        let max_threads = max_threads
            .try_into()
            .map_err(|_| ConfigError::MaxThreads)?;

        let min_threads = min_threads
            .try_into()
            .map_err(|_| ConfigError::MinThreads)?;

        Ok(CpuPool {
            state: Arc::new(CpuPoolState::new(
                name,
                buffer_multiplier,
                min_threads,
                max_threads,
            )),
        })
    }
}

impl Default for Config {
    fn default() -> Self {
        Self::new()
    }
}

/// Errors created by invalid configuration of the CpuPool
#[derive(Debug)]
pub enum ConfigError {
    /// The configured maxumim threads value is lower than the configured minimum threads value
    ThreadCount,

    /// The buffer_multiplier is 0
    BufferMultiplier,

    /// The max_threads value is 0
    MaxThreads,

    /// The min_threads value is 0
    MinThreads,
}

impl std::fmt::Display for ConfigError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::ThreadCount => write!(f, "min_threads cannot be higher than max_threads"),
            Self::BufferMultiplier => write!(f, "buffer_multiplier cannot be zero"),
            Self::MaxThreads => write!(f, "max_threads cannot be zero"),
            Self::MinThreads => write!(f, "min_threads cannot be zero"),
        }
    }
}

impl std::error::Error for ConfigError {}

/// The blocking operation was canceled due to a panic
#[derive(Debug)]
pub struct Canceled;

impl std::fmt::Display for Canceled {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Blocking operation has panicked")
    }
}

impl std::error::Error for Canceled {}

/// The CPUPool handle
#[derive(Clone, Debug)]
pub struct CpuPool {
    state: Arc<CpuPoolState>,
}

impl CpuPool {
    /// Create a new CpuPool with the default configuration
    ///
    /// Example:
    /// ```rust
    /// # use async_cpupool::CpuPool;
    /// let pool = CpuPool::new();
    /// ```
    pub fn new() -> Self {
        Config::default().build().expect("Defaults are valid")
    }

    /// Create a configuration builder to customize the CpuPool
    ///
    /// Example:
    /// ```rust
    /// # use async_cpupool::CpuPool;
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let pool = CpuPool::configure().build()?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn configure() -> Config {
        Config::default()
    }

    /// Spawn a blocking operation on the CpuPool
    ///
    /// Example:
    /// ```rust
    /// # use async_cpupool::CpuPool;
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// # smol::block_on(async {
    /// let pool = CpuPool::new();
    ///
    /// pool.spawn(|| std::thread::sleep(std::time::Duration::from_secs(3))).await?;
    /// # Ok(())
    /// # })
    /// # }
    /// ```
    pub fn spawn<F, T>(&self, send_fn: F) -> impl Future<Output = Result<T, Canceled>> + '_
    where
        F: FnOnce() -> T + Send + 'static,
        T: Send + 'static,
    {
        let (response_tx, response_rx) = spsc::channel();

        let send_fn = Box::new(move || {
            let output = (send_fn)();

            match response_tx.blocking_send(output) {
                Ok(()) => (), // sent
                Err(Canceled) => tracing::warn!("receiver hung up"),
            }
        });

        let opt = self.state.queue.try_push(send_fn);

        let current_threads = self
            .state
            .current_threads
            .load(std::sync::atomic::Ordering::Acquire);

        let pushed = match self.state.queue.is_full_or() {
            Ok(()) => self.push_thread(),
            Err(len) if len > current_threads as usize => self.push_thread(),
            Err(_) => false,
        };

        if pushed {
            tracing::trace!("Pushed thread");
        }

        async {
            if let Some(item) = opt {
                self.state.queue.push(item).await;
            }

            let current_threads = self
                .state
                .current_threads
                .load(std::sync::atomic::Ordering::Acquire);

            match self.state.queue.is_full_or() {
                Ok(()) => {
                    self.push_thread();
                }
                Err(len) if len > current_threads as usize => {
                    self.push_thread();
                }
                Err(len) if len < current_threads.ilog2() as usize => {
                    if let Some(thread) = self.pop_thread() {
                        thread.reap().await;
                    }
                }
                Err(_) => {}
            }

            response_rx.recv().await
        }
    }

    /// Attempt to close the CpuPool
    ///
    /// This operation returns `true` when the pool was succesfully closed, or `false` if there
    /// exist other references to the pool, preventing closure.
    ///
    /// It is not required to call close to close a CpuPool. CpuPools will automatically close
    /// themselves when all clones are dropped. This is simply a method to integrate better with
    /// async runtimes.
    /// Example:
    /// ```rust
    /// # use async_cpupool::CpuPool;
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// # smol::block_on(async {
    /// let pool = CpuPool::new();
    ///
    /// let closed = pool.close().await;
    /// assert!(closed);
    /// # Ok(())
    /// # })
    /// # }
    /// ```
    pub async fn close(self) -> bool {
        let Some(mut state) = Arc::into_inner(self.state) else {
            return false;
        };

        let mut threads = state.take_threads();

        for thread in &mut threads {
            thread.signal.take();
        }

        for mut thread in threads {
            thread.closed.listen().await;

            if let Some(handle) = thread.handle.take() {
                handle.join().expect("Thread panicked");
            }
        }

        true
    }

    fn push_thread(&self) -> bool {
        let current_threads = self
            .state
            .current_threads
            .load(std::sync::atomic::Ordering::Acquire);

        if current_threads >= u64::from(u16::from(self.state.max_threads)) {
            tracing::trace!("At thread maximum");

            return false;
        }

        if self
            .state
            .current_threads
            .compare_exchange(
                current_threads,
                current_threads + 1,
                std::sync::atomic::Ordering::AcqRel,
                std::sync::atomic::Ordering::Relaxed,
            )
            .is_err()
        {
            tracing::trace!("Didn't acquire spawn authorization");

            return false;
        }

        // we updated the count, so we have authorization to spawn a new thread

        let thread_id = self
            .state
            .thread_id
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);

        let thread = spawn(self.state.name, thread_id, self.state.queue.clone());

        self.state
            .threads
            .lock()
            .expect("threads lock poison")
            .push(thread);

        true
    }

    fn pop_thread(&self) -> Option<Thread> {
        let current_threads = self
            .state
            .current_threads
            .load(std::sync::atomic::Ordering::Acquire);

        if current_threads <= u64::from(u16::from(self.state.min_threads)) {
            tracing::info!("At thread minimum");

            return None;
        }

        if self
            .state
            .current_threads
            .compare_exchange(
                current_threads,
                current_threads - 1,
                std::sync::atomic::Ordering::AcqRel,
                std::sync::atomic::Ordering::Relaxed,
            )
            .is_err()
        {
            tracing::trace!("Didn't acquire reap authorization");

            return None;
        }

        // we updated the count, so we have authorization to reap a thread

        self.state
            .threads
            .lock()
            .expect("threads lock poison")
            .pop()
    }
}

impl Default for CpuPool {
    fn default() -> Self {
        Self::new()
    }
}

type SendFn = Box<dyn FnOnce() + Send>;

struct CpuPoolState {
    name: &'static str,
    min_threads: NonZeroU16,
    max_threads: NonZeroU16,
    current_threads: AtomicU64,
    thread_id: AtomicU64,
    queue: Queue<SendFn>,
    threads: Mutex<ThreadVec>,
}

impl CpuPoolState {
    fn new(
        name: &'static str,
        buffer_multiplier: NonZeroUsize,
        min_threads: NonZeroU16,
        max_threads: NonZeroU16,
    ) -> Self {
        let thread_capacity = usize::from(u16::from(max_threads));

        let queue = queue::bounded(usize::from(buffer_multiplier).saturating_mul(thread_capacity));

        let start_threads = u64::from(u16::from(min_threads));

        let threads = ThreadVec::new(start_threads, thread_capacity, |i| {
            spawn(name, i, queue.clone())
        });

        let current_threads = AtomicU64::new(start_threads);
        let thread_id = AtomicU64::new(start_threads);

        CpuPoolState {
            name,
            min_threads,
            max_threads,
            current_threads,
            thread_id,
            queue,
            threads: Mutex::new(threads),
        }
    }

    fn take_threads(&mut self) -> Vec<Thread> {
        self.threads.lock().expect("threads lock poison").take()
    }
}

impl std::fmt::Debug for CpuPoolState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CpuPoolState")
            .field("name", &self.name)
            .field("min_threads", &self.min_threads)
            .field("max_threads", &self.max_threads)
            .finish()
    }
}

struct ThreadVec {
    threads: Vec<Thread>,
}

impl ThreadVec {
    fn new<F>(start_threads: u64, max_threads: usize, spawn: F) -> Self
    where
        F: Fn(u64) -> Thread,
    {
        let mut threads = Vec::with_capacity(max_threads);

        for i in 0..start_threads {
            threads.push((spawn)(i));
        }

        Self { threads }
    }

    fn push(&mut self, thread: Thread) {
        self.threads.push(thread);
    }

    fn pop(&mut self) -> Option<Thread> {
        self.threads.pop()
    }

    fn take(&mut self) -> Vec<Thread> {
        std::mem::take(&mut self.threads)
    }
}

impl Drop for ThreadVec {
    fn drop(&mut self) {
        for thread in &mut self.threads {
            thread.signal.take();
        }

        for thread in &mut self.threads {
            if let Some(handle) = thread.handle.take() {
                handle.join().expect("Thread panicked");
            }
        }
    }
}

struct Thread {
    handle: Option<JoinHandle<()>>,
    signal: Option<DropNotifier>,
    closed: DropListener,
}

impl Thread {
    async fn reap(mut self) {
        self.signal.take();

        self.closed.listen().await;

        if let Some(handle) = self.handle.take() {
            handle.join().expect("Thread panicked");
        }
    }
}

fn spawn(name: &'static str, id: u64, receiver: Queue<SendFn>) -> Thread {
    let (closed_notifier, closed_listener) = drop_notifier::notifier();
    let (signal_notifier, signal_listener) = drop_notifier::notifier();

    let handle = std::thread::Builder::new()
        .name(format!("{name}-{id}"))
        .spawn(move || run(name, id, receiver, signal_listener, closed_notifier))
        .expect("Failed to spawn new thread");

    Thread {
        handle: Some(handle),
        signal: Some(signal_notifier),
        closed: closed_listener,
    }
}

struct MetricsGuard {
    name: &'static str,
    id: u64,
    start: Instant,
    armed: bool,
}

impl MetricsGuard {
    fn guard(name: &'static str, id: u64) -> Self {
        tracing::trace!("Starting {name}-{id}");
        metrics::counter!(format!("async-cpupool.{name}.thread.launched")).increment(1);

        MetricsGuard {
            name,
            id,
            start: Instant::now(),
            armed: true,
        }
    }

    fn disarm(mut self) {
        self.armed = false;
    }
}

impl Drop for MetricsGuard {
    fn drop(&mut self) {
        metrics::counter!(format!("async-cpupool.{}.thread.closed", self.name), "clean" => (!self.armed).to_string()).increment(1);
        metrics::histogram!(format!("async-cpupool.{}.thread.seconds", self.name), "clean" => (!self.armed).to_string()).record(self.start.elapsed().as_secs_f64());
        tracing::trace!("Stopping {}-{}", self.name, self.id);
    }
}

fn run(
    name: &'static str,
    id: u64,
    receiver: Queue<SendFn>,
    signal: DropListener,
    closed_tx: DropNotifier,
) {
    let guard = MetricsGuard::guard(name, id);

    let mut signal = std::pin::pin!(signal.listen());

    loop {
        match block_on(select(&mut signal, receiver.pop())) {
            selector::Either::Left(_) => break,
            selector::Either::Right(send_fn) => invoke_send_fn(name, send_fn),
        }
    }

    guard.disarm();

    drop(closed_tx);
}

fn invoke_send_fn(name: &'static str, send_fn: SendFn) {
    let start = Instant::now();
    metrics::counter!(format!("async-cpupool.{name}.operation.start")).increment(1);

    let res = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || {
        (send_fn)();
    }));

    metrics::counter!(format!("async-cpupool.{name}.operation.end"), "complete" => res.is_ok().to_string()).increment(1);
    metrics::histogram!(format!("async-cpupool.{name}.operation.seconds"), "complete" => res.is_ok().to_string()).record(start.elapsed().as_secs_f64());

    if let Err(e) = res {
        tracing::trace!("panic in spawned task: {e:?}");
    }
}