use std::thread::JoinHandle;
use std::time::Duration;

use async_channel::{Receiver, Sender};
use rusqlite::Connection;

pub type RusqliteThreadPoolTask =
	Box<dyn FnOnce(&mut Connection) -> () + Send + Sync + 'static>;
pub type RusqliteThreadPoolWorker =
	(JoinHandle<()>, Sender<RusqliteThreadPoolTask>);

pub type PoolReceiver = Receiver<RusqliteThreadPoolTask>;

#[derive(Clone)]
pub struct PoolSender {
	tx: Sender<RusqliteThreadPoolTask>,
}
impl PoolSender {
	pub async fn send<T>(
		&self,
		f: impl FnOnce(&mut Connection) -> rusqlite::Result<T>
			+ Send
			+ Sync
			+ 'static,
	) -> rusqlite::Result<T>
	where
		T: Send + Sync + 'static,
	{
		let (tx, rx) = async_channel::bounded(1);
		let task = move |conn: &mut Connection| {
			let out = f(conn);
			tx.send_blocking(out).unwrap();
		};
		let task: RusqliteThreadPoolTask = Box::new(task);
		self.tx.send(task).await.unwrap();

		rx.recv().await.unwrap()
	}
}

pub fn spawn_worker(
	mut conn: Connection,
	rx: Receiver<RusqliteThreadPoolTask>,
) -> JoinHandle<()> {
	std::thread::spawn(move || {
		while let Ok(task) = rx.try_recv() {
			task(&mut conn);
		}
	})
}

pub fn spawn_threadpool(
	n_workers: u8,
	db_path: &str,
	mut init: impl FnMut(&mut Connection) -> rusqlite::Result<()>,
) -> rusqlite::Result<(PoolSender, Vec<JoinHandle<()>>, PoolReceiver)> {
	assert!(n_workers > 0);

	let (tx, rx) = async_channel::unbounded();
	let mut workers = vec![];
	for _ in 0..n_workers {
		let mut conn = Connection::open(db_path)?;
		init(&mut conn)?;

		let handle = spawn_worker(conn, rx.clone());

		workers.push(handle);
	}

	return Ok((PoolSender { tx }, workers, rx));
}

pub fn spawn_threadpool_supervised(
	n_workers: u8,
	db_path: &str,
	mut init: impl FnMut(&mut Connection) -> rusqlite::Result<()>
		+ Send
		+ Sync
		+ Clone
		+ 'static,
) -> rusqlite::Result<(PoolSender, JoinHandle<()>, PoolReceiver)> {
	let (tx, mut workers, rx) =
		spawn_threadpool(n_workers, db_path, init.clone())?;

	let rxt = rx.clone();
	let db_path = db_path.to_string();
	let supervisor = std::thread::spawn(move || {
		let rx = rxt;
		let db_path = db_path;

		loop {
			let mut dead_worker_idxs = vec![];
			for (i, worker) in workers.iter().enumerate() {
				if worker.is_finished() {
					dead_worker_idxs.push(i);
				}
			}

			for i in dead_worker_idxs {
				let mut conn = Connection::open(&db_path).unwrap();
				init(&mut conn).unwrap();

				workers[i] = spawn_worker(conn, rx.clone());
			}

			std::thread::sleep(Duration::from_secs(1));
		}
	});

	return Ok((tx, supervisor, rx));
}
