use std::marker::PhantomData;
use std::sync::Arc;
use std::sync::mpsc::{self, Receiver, RecvTimeoutError, SyncSender, TrySendError};
use std::thread::{self, JoinHandle, Scope};
use std::time::Duration;

/// Executor shared by every worker of a pool. The `'static` bound required to
/// move the executor into detached threads is expressed on the pool types
/// instead of this trait so scoped pools can borrow their environment.
pub trait ObjectTaskExecutor<T, R>: Send + Sync {
    fn execute(&self, worker_index: usize, task: T) -> R;
}

enum ObjectWorkerMessage<T> {
    Task(T),
    Shutdown,
}

#[derive(Debug)]
pub enum ObjectWorkerSubmitError<T> {
    QueueFull { worker_index: usize, task: T },
    Disconnected { worker_index: usize, task: T },
}

pub struct ObjectWorkerPool<T, R, E>
where
    T: Send + 'static,
    R: Send + 'static,
    E: ObjectTaskExecutor<T, R> + 'static,
{
    task_txs: Vec<SyncSender<ObjectWorkerMessage<T>>>,
    result_rx: Receiver<R>,
    workers: Vec<JoinHandle<()>>,
    next_worker_idx: usize,
    _executor: Arc<E>,
}

impl<T, R, E> ObjectWorkerPool<T, R, E>
where
    T: Send + 'static,
    R: Send + 'static,
    E: ObjectTaskExecutor<T, R> + 'static,
{
    pub fn new(worker_count: usize, queue_capacity: usize, executor: E) -> Result<Self, String> {
        if worker_count == 0 {
            return Err("ObjectWorkerPool requires at least one worker".to_string());
        }
        if queue_capacity == 0 {
            return Err("ObjectWorkerPool requires queue_capacity > 0".to_string());
        }

        let executor = Arc::new(executor);
        let (result_tx, result_rx) = mpsc::channel::<R>();
        let mut task_txs = Vec::with_capacity(worker_count);
        let mut workers = Vec::with_capacity(worker_count);

        for worker_index in 0..worker_count {
            let (task_tx, task_rx) = mpsc::sync_channel::<ObjectWorkerMessage<T>>(queue_capacity);
            let result_tx = result_tx.clone();
            let executor = Arc::clone(&executor);
            let handle = thread::Builder::new()
                .name(format!("object-validation-worker-{worker_index}"))
                .spawn(move || object_worker_loop(worker_index, task_rx, result_tx, executor))
                .map_err(|e| format!("spawn object worker failed: {e}"))?;
            task_txs.push(task_tx);
            workers.push(handle);
        }

        Ok(Self {
            task_txs,
            result_rx,
            workers,
            next_worker_idx: 0,
            _executor: executor,
        })
    }

    pub fn worker_count(&self) -> usize {
        self.task_txs.len()
    }

    pub fn next_worker_index(&self) -> usize {
        self.next_worker_idx
    }

    pub fn try_submit_round_robin(&mut self, task: T) -> Result<usize, ObjectWorkerSubmitError<T>> {
        let worker_index = self.next_worker_idx % self.task_txs.len();
        match self.task_txs[worker_index].try_send(ObjectWorkerMessage::Task(task)) {
            Ok(()) => {
                self.next_worker_idx = (worker_index + 1) % self.task_txs.len();
                Ok(worker_index)
            }
            Err(TrySendError::Full(ObjectWorkerMessage::Task(task))) => {
                Err(ObjectWorkerSubmitError::QueueFull { worker_index, task })
            }
            Err(TrySendError::Disconnected(ObjectWorkerMessage::Task(task))) => {
                Err(ObjectWorkerSubmitError::Disconnected { worker_index, task })
            }
            Err(TrySendError::Full(ObjectWorkerMessage::Shutdown))
            | Err(TrySendError::Disconnected(ObjectWorkerMessage::Shutdown)) => {
                unreachable!("shutdown is never submitted via try_submit_round_robin")
            }
        }
    }

    pub fn recv_result_timeout(&self, timeout: Duration) -> Result<Option<R>, String> {
        match self.result_rx.recv_timeout(timeout) {
            Ok(result) => Ok(Some(result)),
            Err(RecvTimeoutError::Timeout) => Ok(None),
            Err(RecvTimeoutError::Disconnected) => {
                Err("object worker result channel disconnected".to_string())
            }
        }
    }

    pub fn shutdown(mut self) -> Result<(), String> {
        self.shutdown_inner()
    }

    fn shutdown_inner(&mut self) -> Result<(), String> {
        if self.workers.is_empty() {
            return Ok(());
        }
        for tx in &self.task_txs {
            tx.send(ObjectWorkerMessage::Shutdown)
                .map_err(|e| format!("send shutdown to object worker failed: {e}"))?;
        }
        let mut first_err = None;
        for handle in self.workers.drain(..) {
            if let Err(e) = handle.join() {
                if first_err.is_none() {
                    first_err = Some(format!("join object worker failed: {e:?}"));
                }
            }
        }
        if let Some(err) = first_err {
            return Err(err);
        }
        Ok(())
    }
}

impl<T, R, E> Drop for ObjectWorkerPool<T, R, E>
where
    T: Send + 'static,
    R: Send + 'static,
    E: ObjectTaskExecutor<T, R> + 'static,
{
    fn drop(&mut self) {
        let _ = self.shutdown_inner();
    }
}

fn object_worker_loop<T, R, E>(
    worker_index: usize,
    task_rx: Receiver<ObjectWorkerMessage<T>>,
    result_tx: mpsc::Sender<R>,
    executor: Arc<E>,
) where
    T: Send,
    R: Send,
    E: ObjectTaskExecutor<T, R>,
{
    loop {
        match task_rx.recv() {
            Ok(ObjectWorkerMessage::Task(task)) => {
                let result = executor.execute(worker_index, task);
                if result_tx.send(result).is_err() {
                    break;
                }
            }
            Ok(ObjectWorkerMessage::Shutdown) | Err(_) => break,
        }
    }
}

/// Scoped variant of `ObjectWorkerPool`: workers are spawned on a
/// `std::thread::Scope`, so tasks, results and the executor may borrow their
/// environment (`'env`) instead of being `'static`. The pool never sends
/// `Shutdown`; workers exit when every task sender is dropped, which happens
/// when the pool itself is dropped ahead of the scope join.
pub struct ScopedObjectWorkerPool<'scope, 'env, T, R, E>
where
    T: Send + 'env,
    R: Send + 'env,
    E: ObjectTaskExecutor<T, R> + 'env,
{
    task_txs: Vec<SyncSender<ObjectWorkerMessage<T>>>,
    result_rx: Receiver<R>,
    next_worker_idx: usize,
    _executor: Arc<E>,
    // Join handles are intentionally not stored: dropping a `ScopedJoinHandle`
    // detaches the worker and the enclosing scope joins it on exit, after the
    // dropped task senders have made every worker return.
    _marker: PhantomData<(&'scope (), &'env ())>,
}

impl<'scope, 'env, T, R, E> ScopedObjectWorkerPool<'scope, 'env, T, R, E>
where
    T: Send + 'env,
    R: Send + 'env,
    E: ObjectTaskExecutor<T, R> + 'env,
{
    pub fn new(
        scope: &'scope Scope<'scope, 'env>,
        worker_count: usize,
        queue_capacity: usize,
        executor: E,
    ) -> Result<Self, String> {
        if worker_count == 0 {
            return Err("ScopedObjectWorkerPool requires at least one worker".to_string());
        }
        if queue_capacity == 0 {
            return Err("ScopedObjectWorkerPool requires queue_capacity > 0".to_string());
        }

        let executor = Arc::new(executor);
        let (result_tx, result_rx) = mpsc::channel::<R>();
        let mut task_txs = Vec::with_capacity(worker_count);

        for worker_index in 0..worker_count {
            let (task_tx, task_rx) = mpsc::sync_channel::<ObjectWorkerMessage<T>>(queue_capacity);
            let result_tx = result_tx.clone();
            let executor = Arc::clone(&executor);
            thread::Builder::new()
                .name(format!("object-validation-worker-{worker_index}"))
                .spawn_scoped(scope, move || {
                    object_worker_loop(worker_index, task_rx, result_tx, executor)
                })
                .map_err(|e| format!("spawn scoped object worker failed: {e}"))?;
            task_txs.push(task_tx);
        }

        Ok(Self {
            task_txs,
            result_rx,
            next_worker_idx: 0,
            _executor: executor,
            _marker: PhantomData,
        })
    }

    pub fn worker_count(&self) -> usize {
        self.task_txs.len()
    }

    pub fn try_submit_round_robin(&mut self, task: T) -> Result<usize, ObjectWorkerSubmitError<T>> {
        let worker_index = self.next_worker_idx % self.task_txs.len();
        match self.task_txs[worker_index].try_send(ObjectWorkerMessage::Task(task)) {
            Ok(()) => {
                self.next_worker_idx = (worker_index + 1) % self.task_txs.len();
                Ok(worker_index)
            }
            Err(TrySendError::Full(ObjectWorkerMessage::Task(task))) => {
                Err(ObjectWorkerSubmitError::QueueFull { worker_index, task })
            }
            Err(TrySendError::Disconnected(ObjectWorkerMessage::Task(task))) => {
                Err(ObjectWorkerSubmitError::Disconnected { worker_index, task })
            }
            Err(TrySendError::Full(ObjectWorkerMessage::Shutdown))
            | Err(TrySendError::Disconnected(ObjectWorkerMessage::Shutdown)) => {
                unreachable!("shutdown is never submitted via try_submit_round_robin")
            }
        }
    }

    pub fn recv_result_timeout(&self, timeout: Duration) -> Result<Option<R>, String> {
        match self.result_rx.recv_timeout(timeout) {
            Ok(result) => Ok(Some(result)),
            Err(RecvTimeoutError::Timeout) => Ok(None),
            Err(RecvTimeoutError::Disconnected) => {
                Err("scoped object worker result channel disconnected".to_string())
            }
        }
    }
}

impl<'scope, 'env, T, R, E> Drop for ScopedObjectWorkerPool<'scope, 'env, T, R, E>
where
    T: Send + 'env,
    R: Send + 'env,
    E: ObjectTaskExecutor<T, R> + 'env,
{
    fn drop(&mut self) {
        // Close every worker input queue so blocked `recv` calls return and
        // the scoped workers exit before the enclosing scope joins them.
        self.task_txs.clear();
    }
}

#[cfg(test)]
mod tests {
    use super::{ObjectTaskExecutor, ObjectWorkerPool, ObjectWorkerSubmitError};
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::sync::{Arc, Barrier};
    use std::time::Duration;

    #[derive(Clone)]
    struct EchoExecutor;

    impl ObjectTaskExecutor<u32, (usize, u32)> for EchoExecutor {
        fn execute(&self, worker_index: usize, task: u32) -> (usize, u32) {
            (worker_index, task)
        }
    }

    #[test]
    fn object_worker_pool_rejects_invalid_config_and_shutdowns_explicitly() {
        let err = match ObjectWorkerPool::new(0, 1, EchoExecutor) {
            Ok(_) => panic!("zero workers should be rejected"),
            Err(err) => err,
        };
        assert!(err.contains("at least one worker"));
        let err = match ObjectWorkerPool::new(1, 0, EchoExecutor) {
            Ok(_) => panic!("zero queue should be rejected"),
            Err(err) => err,
        };
        assert!(err.contains("queue_capacity > 0"));

        let pool = ObjectWorkerPool::new(2, 1, EchoExecutor).expect("pool");
        assert_eq!(pool.worker_count(), 2);
        assert_eq!(pool.next_worker_index(), 0);
        pool.shutdown().expect("shutdown");
    }

    #[test]
    fn object_worker_pool_round_robin_submits_to_worker_queues() {
        let mut pool = ObjectWorkerPool::new(3, 4, EchoExecutor).expect("pool");
        assert_eq!(pool.try_submit_round_robin(10).expect("submit 10"), 0);
        assert_eq!(pool.try_submit_round_robin(11).expect("submit 11"), 1);
        assert_eq!(pool.try_submit_round_robin(12).expect("submit 12"), 2);
        assert_eq!(pool.try_submit_round_robin(13).expect("submit 13"), 0);

        let mut results = Vec::new();
        for _ in 0..4 {
            results.push(
                pool.recv_result_timeout(Duration::from_secs(1))
                    .expect("result channel")
                    .expect("result"),
            );
        }
        results.sort_by_key(|(_, task)| *task);
        assert_eq!(results, vec![(0, 10), (1, 11), (2, 12), (0, 13)]);
    }

    struct BlockingExecutor {
        barrier: Arc<Barrier>,
        started: Arc<AtomicBool>,
    }

    impl ObjectTaskExecutor<u32, u32> for BlockingExecutor {
        fn execute(&self, _worker_index: usize, task: u32) -> u32 {
            self.started.store(true, Ordering::SeqCst);
            self.barrier.wait();
            task
        }
    }

    #[test]
    fn object_worker_pool_reports_full_worker_queue_without_advancing_round_robin() {
        let barrier = Arc::new(Barrier::new(2));
        let started = Arc::new(AtomicBool::new(false));
        let mut pool = ObjectWorkerPool::new(
            1,
            1,
            BlockingExecutor {
                barrier: Arc::clone(&barrier),
                started: Arc::clone(&started),
            },
        )
        .expect("pool");

        assert_eq!(pool.try_submit_round_robin(1).expect("first task"), 0);
        let deadline = std::time::Instant::now() + Duration::from_secs(1);
        while !started.load(Ordering::SeqCst) {
            assert!(
                std::time::Instant::now() < deadline,
                "worker did not start first task"
            );
            std::thread::sleep(Duration::from_millis(1));
        }
        assert_eq!(pool.try_submit_round_robin(2).expect("queued task"), 0);
        match pool.try_submit_round_robin(3) {
            Err(ObjectWorkerSubmitError::QueueFull { worker_index, task }) => {
                assert_eq!(worker_index, 0);
                assert_eq!(task, 3);
            }
            other => panic!("expected queue full, got {other:?}"),
        }
        assert_eq!(pool.next_worker_index(), 0);

        barrier.wait();
        assert_eq!(
            pool.recv_result_timeout(Duration::from_secs(1))
                .expect("result channel"),
            Some(1)
        );
        barrier.wait();
        assert_eq!(
            pool.recv_result_timeout(Duration::from_secs(1))
                .expect("result channel"),
            Some(2)
        );
    }

    struct BorrowingEchoExecutor<'a> {
        base: &'a u32,
    }

    impl<'a> ObjectTaskExecutor<u32, u32> for BorrowingEchoExecutor<'a> {
        fn execute(&self, _worker_index: usize, task: u32) -> u32 {
            task + *self.base
        }
    }

    #[test]
    fn scoped_object_worker_pool_borrows_environment_and_processes_tasks() {
        let base = 100u32;
        std::thread::scope(|scope| {
            let mut pool = super::ScopedObjectWorkerPool::new(
                scope,
                2,
                2,
                BorrowingEchoExecutor { base: &base },
            )
            .expect("scoped pool");
            assert_eq!(pool.worker_count(), 2);
            pool.try_submit_round_robin(1).expect("submit 1");
            pool.try_submit_round_robin(2).expect("submit 2");
            let mut results = Vec::new();
            for _ in 0..2 {
                results.push(
                    pool.recv_result_timeout(Duration::from_secs(1))
                        .expect("result channel")
                        .expect("result"),
                );
            }
            results.sort();
            assert_eq!(results, vec![101, 102]);
            // Dropping the pool inside the scope closes the task queues; the
            // workers exit on their own and the scope join below must not hang.
            drop(pool);
        });
    }

    #[test]
    fn scoped_object_worker_pool_reports_full_queue() {
        std::thread::scope(|scope| {
            let barrier = Arc::new(Barrier::new(2));
            let started = Arc::new(AtomicBool::new(false));
            let mut pool = super::ScopedObjectWorkerPool::new(
                scope,
                1,
                1,
                BlockingExecutor {
                    barrier: Arc::clone(&barrier),
                    started: Arc::clone(&started),
                },
            )
            .expect("scoped pool");
            pool.try_submit_round_robin(1).expect("first task");
            let deadline = std::time::Instant::now() + Duration::from_secs(1);
            while !started.load(Ordering::SeqCst) {
                assert!(
                    std::time::Instant::now() < deadline,
                    "scoped worker did not start first task"
                );
                std::thread::sleep(Duration::from_millis(1));
            }
            pool.try_submit_round_robin(2).expect("queued task");
            match pool.try_submit_round_robin(3) {
                Err(ObjectWorkerSubmitError::QueueFull { worker_index, task }) => {
                    assert_eq!(worker_index, 0);
                    assert_eq!(task, 3);
                }
                other => panic!("expected queue full, got {other:?}"),
            }
            // Dropping the pool closes the task queues; releasing the barrier
            // afterwards lets the blocked worker finish task 1, observe the
            // closed result channel, and exit before the scope join.
            drop(pool);
            barrier.wait();
        });
    }
}