use std::cell::RefCell;
use std::sync::{Arc, Barrier, Mutex};
use std::thread::{Builder, JoinHandle, LocalKey};
use std::time::Duration;

use anyhow::anyhow;

use crate::blocking_queue::BlockingQueue;
use crate::command::Command;
use crate::shutdown_mode::ShutdownMode;
use crate::signal::Signal;

struct RunInAllThreadsCommand {
    f: Arc<dyn Fn() + Send + Sync>,
    b: Arc<Barrier>,
}

impl RunInAllThreadsCommand {
    pub fn new(f: Arc<dyn Fn() + Send + Sync>, b: Arc<Barrier>) -> RunInAllThreadsCommand {
        RunInAllThreadsCommand {
            f,
            b,
        }
    }
}

impl Command for RunInAllThreadsCommand {
    fn execute(&self) -> Result<(), anyhow::Error> {
        {
            (self.f)();
        }
        self.b.wait();
        Ok(())
    }
}

struct RunMutInAllThreadsCommand {
    f: Arc<Mutex<dyn FnMut() + Send + Sync>>,
    b: Arc<Barrier>,
}

impl RunMutInAllThreadsCommand {
    pub fn new(f: Arc<Mutex<dyn FnMut() + Send + Sync>>, b: Arc<Barrier>) -> RunMutInAllThreadsCommand {
        RunMutInAllThreadsCommand {
            f,
            b,
        }
    }
}

impl Command for RunMutInAllThreadsCommand {
    fn execute(&self) -> Result<(), anyhow::Error> {
        {
            let mut f = self.f.lock().unwrap();
            f();
        }
        self.b.wait();
        Ok(())
    }
}

/// Execute tasks concurrently while maintaining bounds on memory consumption
///
/// To demonstrate the use case this implementation solves let's consider a program that reads
/// lines from a file and writes those lines to another file after some processing. The processing
/// itself is stateless and can be done in parallel on each line, but the reading and writing must
/// be sequential. Using this implementation we will read the input in the main thread, submit it
/// for concurrent processing to a processing thread pool and collect it for writing in a writing thread pool
/// with a single thread. See ./examples/read_process_write_pipeline.rs. The submission to a thread
/// pool is done through a blocking bounded queue, so if the processing thread pool or the writing
/// thread pool cannot keep up, their blocking queues will fill up and create a backpressure that
/// will pause the reading. So the resulting pipeline will stabilize on a throughput commanded by the
/// slowest stage with the memory consumption determined by sizes of queues and number of
/// threads in each thread pool.
///
/// For reference see [Command Pattern](https://en.wikipedia.org/wiki/Command_pattern) and
/// [Producer-Consumer](https://en.wikipedia.org/wiki/Producer%E2%80%93consumer_problem)
///
pub struct ThreadPool {
    name: String,
    tasks: usize,
    queue: Arc<BlockingQueue<Box<dyn Command + Send + Sync>, Signal>>,
    threads: Vec<JoinHandle<Result<(), anyhow::Error>>>,
    join_error_handler: fn(String, String),
    shutdown_mode: ShutdownMode,
    expired: bool,
}

impl ThreadPool {
    pub(crate) fn new(
        name: String,
        tasks: usize,
        queue_size: usize,
        join_error_handler: fn(String, String),
        shutdown_mode: ShutdownMode,
    ) -> Result<ThreadPool, anyhow::Error> {
        let start_barrier = Arc::new(Barrier::new(tasks + 1));
        let mut threads = Vec::<JoinHandle<Result<(), anyhow::Error>>>::new();
        let queue = Arc::new(
            BlockingQueue::<Box<dyn Command + Send + Sync>, Signal>::new(queue_size)
        );
        for i in 0..tasks {
            let barrier = start_barrier.clone();
            let t = Self::create_thread(
                &name,
                i,
                barrier,
                queue.clone(),
            );
            threads.push(t.unwrap());
        }

        start_barrier.wait();

        Ok(
            ThreadPool {
                name,
                tasks,
                queue: queue.clone(),
                threads,
                join_error_handler,
                shutdown_mode,
                expired: false,
            }
        )
    }

    /// Set the number of concurrent threads in the thread pool
    pub fn tasks(&self) -> usize {
        self.tasks
    }

    fn create_thread(
        name: &String,
        index: usize,
        barrier: Arc<Barrier>,
        queue: Arc<BlockingQueue<Box<dyn Command + Send + Sync>, Signal>>,
    ) -> Result<JoinHandle<Result<(), anyhow::Error>>, anyhow::Error> {
        let builder = Builder::new();
        Ok(builder
            .name(format!("{name}-{index}"))
            .spawn(move || {
                    barrier.wait();
                    let mut r: Result<(), anyhow::Error> = Ok(());
                    loop {
                        let (command, signal) = queue.dequeue();
                        if let Some(c) = command {
                            match c.execute() {
                                Ok(_) => {}
                                Err(e) => {
                                    r = Err(e);
                                }
                            }
                        }
                        if let Some(s) = signal {
                            match s {
                                Signal::Shutdown => {
                                    break r;
                                }
                            }
                        }
                    }
                }
            )?
        )
    }

    /// Execute f in all threads.
    ///
    /// This function returns only after f had completed in all threads. Can be used to collect
    /// data produced by the threads. See ./examples/fetch_thread_local.rs.
    ///
    /// Caveat: this is a [barrier](https://en.wikipedia.org/wiki/Barrier_%28computer_science%29)
    /// function. So if one of the threads is busy with a long running task or is deadlocked, this
    /// will halt all the threads until f can be executed.
    pub fn in_all_threads_mut(&self, f: Arc<Mutex<dyn FnMut() + Send + Sync>>) {
        let b = Arc::new(Barrier::new(self.tasks + 1));
        for _i in 0..self.tasks {
            self.submit(Box::new(RunMutInAllThreadsCommand::new(f.clone(), b.clone())));
        }
        b.wait();
    }

    /// Execute f in all threads.
    ///
    /// This function returns only after f had completed in all threads. Can be used to flush
    /// data produced by the threads or simply execute work concurrently. See ./examples/flush_thread_local.rs.
    ///
    /// Caveat: this is a [barrier](https://en.wikipedia.org/wiki/Barrier_%28computer_science%29)
    /// function. So if one of the threads is busy with a long running task or is deadlocked, this
    /// will halt all the threads until f can be executed.
    pub fn in_all_threads(&self, f: Arc<dyn Fn() + Send + Sync>) {
        let b = Arc::new(Barrier::new(self.tasks + 1));
        for _i in 0..self.tasks {
            self.submit(Box::new(RunInAllThreadsCommand::new(f.clone(), b.clone())));
        }
        b.wait();
    }

    /// Initializes the `local_key` to contain `val`.
    ///
    /// See ./examples/thread_local.rs
    pub fn set_thread_local<T>(&mut self, local_key: &'static LocalKey<RefCell<T>>, val: T)
        where T: Sync + Send + Clone {
        self.in_all_threads_mut(
            Arc::new(
                Mutex::new(
                    move || {
                        local_key.with(
                            |value| {
                                value.replace(val.clone())
                            }
                        );
                    }
                )
            )
        );
    }


    /// Shut down the thread pool.
    ///
    /// This will shut down the thread pool according to configuration. When configured with
    /// * [ShutdownMode::Immediate] - terminate each tread after completing the current task
    /// * [ShutdownMode::CompletePending] - terminate after completing all pending tasks
    pub fn shutdown(&mut self) {
        self.expired = true;
        match self.shutdown_mode {
            ShutdownMode::Immediate => {
                self.queue.signal(Signal::Shutdown);
            }
            ShutdownMode::CompletePending => {
                self.queue.wait_empty(Duration::MAX);
                self.queue.signal(Signal::Shutdown);
            }
        }
    }

    /// Wait until all thread pool threads completed.
    pub fn join(&mut self) -> Result<(), anyhow::Error> {
        let mut join_errors = Vec::<String>::new();
        while self.threads.len() > 0 {
            let t = self.threads.pop().unwrap();
            let name = t.thread().name().unwrap_or("unnamed").to_string();
            match t.join() {
                Ok(r) => {
                    match r {
                        Ok(_) => {}
                        Err(e) => {
                            let message = format!("{e:?}");
                            join_errors.push(message.clone());
                            (self.join_error_handler)(name, message);
                        }
                    }
                }
                Err(e) => {
                    let mut message = "Unknown error".to_string();
                    if let Some(error) = e.downcast_ref::<&'static str>() {
                        message = error.to_string();
                    }
                    join_errors.push(message.clone());
                    (self.join_error_handler)(name, message);
                }
            }
        }
        if join_errors.is_empty() {
            Ok(())
        } else {
            Err(anyhow!("Errors occurred while joining threads in the {} pool: {}", self.name, join_errors.join(", "))
            )
        }
    }

    /// Submit command for execution
    pub fn submit(&self, command: Box<dyn Command + Send + Sync>) {
        if self.expired {
            panic!("the thread pool {} is expired", self.name)
        }
        self.try_submit(command, Duration::MAX);
    }

    /// Submit command for execution with timeout
    ///
    /// Returns the command on failure and None on success
    pub fn try_submit(&self, command: Box<dyn Command + Send + Sync>, timeout: Duration) -> Option<Box<dyn Command + Send + Sync>> {
        self.queue.try_enqueue(command, timeout)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::{AtomicUsize, Ordering};

    use crate::shutdown_mode::ShutdownMode;
    use crate::shutdown_mode::ShutdownMode::CompletePending;
    use crate::thread_pool_builder::ThreadPoolBuilder;

    use super::*;

    struct TestCommand {
        _payload: i32,
        execution_counter: Arc<AtomicUsize>,
    }

    impl TestCommand {
        pub fn new(payload: i32, execution_counter: Arc<AtomicUsize>) -> TestCommand {
            TestCommand {
                _payload: payload,
                execution_counter,
            }
        }
    }

    impl Command for TestCommand {
        fn execute(&self) -> Result<(), anyhow::Error> {
            self.execution_counter.fetch_add(1, Ordering::Relaxed);
            Ok(())
        }
    }

    #[test]
    fn test_create() {
        let mut thread_pool_builder = ThreadPoolBuilder::new();
        let tp_result = thread_pool_builder
            .name("t".to_string())
            .tasks(4)
            .queue_size(8)
            .build();

        match tp_result {
            Ok(mut tp) => {
                assert!(true);
                tp.shutdown();
                assert_eq!((), tp.join().unwrap());
            }
            Err(_) => {
                assert!(false);
            }
        }
    }

    #[test]
    fn test_submit() {
        let mut thread_pool_builder = ThreadPoolBuilder::new();
        let mut tp = thread_pool_builder
            .name("t".to_string())
            .tasks(4)
            .queue_size(2048)
            .build()
            .unwrap();

        let execution_counter = Arc::new(AtomicUsize::from(0));
        for _i in 0..1024 {
            let ec = execution_counter.clone();
            tp.submit(Box::new(TestCommand::new(4, ec)));
        }

        tp.shutdown();
        tp.join().expect("Failed to join thread pool");
        assert_eq!((), tp.join().unwrap());
        // accidental but usually works
        // if fails safe to comment out the next two lines
        // assert!(execution_counter.fetch_or(0, Ordering::Relaxed) > 0);
        // assert!(execution_counter.fetch_or(0, Ordering::Relaxed) < 1024);
    }

    #[test]
    fn test_shutdown_complete_pending() {
        let mut thread_pool_builder = ThreadPoolBuilder::new();
        let mut tp = thread_pool_builder
            .name("t".to_string())
            .tasks(4)
            .queue_size(2048)
            .shutdown_mode(ShutdownMode::CompletePending)
            .build()
            .unwrap();

        let execution_counter = Arc::new(AtomicUsize::from(0));
        for _i in 0..1024 {
            let ec = execution_counter.clone();
            tp.submit(Box::new(TestCommand::new(4, ec)));
        }

        tp.shutdown();
        tp.join().expect("Failed to join thread pool");
        assert_eq!((), tp.join().unwrap());
        assert_eq!(execution_counter.fetch_or(0, Ordering::Relaxed), 1024);
    }

    struct PanicTestCommand {}

    impl PanicTestCommand {
        pub fn new() -> PanicTestCommand {
            PanicTestCommand {}
        }
    }

    impl Command for PanicTestCommand {
        fn execute(&self) -> Result<(), anyhow::Error> {
            Err(anyhow!("simulating error during command execution"))
        }
    }

    #[test]
    fn test_join_error_handler() {
        let mut thread_pool_builder = ThreadPoolBuilder::new();
        let mut tp = thread_pool_builder
            .name("t".to_string())
            .tasks(4)
            .shutdown_mode(CompletePending)
            .queue_size(8)
            .join_error_handler(
                |name, message| {
                    println!("Thread {name} ended with and error {message}")
                }
            )
            .build()
            .unwrap();

        for _i in 0..2 {
            tp.submit(Box::new(PanicTestCommand::new()));
        }

        tp.shutdown();
        let r = tp.join();
        assert!(r.is_err());
    }

    #[test]
    #[should_panic]
    fn test_use_after_join() {
        let mut thread_pool_builder = ThreadPoolBuilder::new();
        let mut tp = thread_pool_builder
            .name("t".to_string())
            .tasks(4)
            .queue_size(2048)
            .shutdown_mode(ShutdownMode::CompletePending)
            .build()
            .unwrap();

        let execution_counter = Arc::new(AtomicUsize::from(0));
        for _i in 0..1024 {
            let ec = execution_counter.clone();
            tp.submit(Box::new(TestCommand::new(4, ec)));
        }

        tp.shutdown();
        tp.join().expect("Failed to join thread pool");
        let execution_counter = Arc::new(AtomicUsize::from(0));
        for _i in 0..1024 {
            let ec = execution_counter.clone();
            tp.submit(Box::new(TestCommand::new(4, ec)));
        }
    }
}