use std::future::Future;

use tokio::task::JoinHandle;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum RuntimeMode {
    /// Spawns tasks on the current runtime that has invoked [`ServiceBuilder::build`] or
    /// [`ServiceBuilder::build_mutable`].
    ///
    /// This is the default.
    #[default]
    Current,

    /// Spawns a single threaded [`tokio::runtime::Runtime`] for each shard.
    ///
    /// This is useful when you are running into limitations with having a single runtime for all the
    /// shards, because currently, a single tokio runtime becomes in-efficient when it spans a large number of worker
    /// threads across multiple NUMA nodes.
    ///
    /// See [this issue](https://github.com/tokio-rs/tokio/issues/5076) for more information.
    RuntimePerShard,
}

impl RuntimeMode {
    pub(super) fn spawn_shard(
        &self,
        shard_id: u8,
        shard_future: impl Future<Output = ()> + Send + 'static,
    ) -> JoinHandle<()> {
        match self {
            Self::Current => tokio::task::spawn(shard_future),
            Self::RuntimePerShard => {
                let runtime = tokio::runtime::Builder::new_current_thread()
                    .thread_name_fn(move || format!("shard-{}-runtime-thread", shard_id))
                    .enable_all()
                    .build()
                    .expect("failed to build multi thread runtime");

                let (shard_finished_tx, shard_finished_rx) = tokio::sync::oneshot::channel();
                let join_handle = runtime.spawn(async move {
                    shard_future.await;
                    shard_finished_tx.send(()).ok();
                });

                std::thread::Builder::new()
                    .name(format!("shard-{}-runtime-main-thread", shard_id))
                    .spawn(move || runtime.block_on(shard_finished_rx))
                    .unwrap_or_else(|_| {
                        panic!("failed to spawn runtime main thread for shard={}", shard_id)
                    });

                join_handle
            }
        }
    }
}
