use crate::background_job::DEFAULT_QUEUE;
use crate::job_registry::JobRegistry;
use crate::worker::Worker;
use crate::{BackgroundJob, listener, storage};
use anyhow::anyhow;
use diesel_async::AsyncPgConnection;
use diesel_async::pooled_connection::deadpool::Pool;
use futures_util::FutureExt;
use futures_util::future::join_all;
use std::collections::HashMap;
use std::panic::AssertUnwindSafe;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Notify;
use tokio::task::JoinHandle;
use tracing::{Instrument, error, info, info_span, warn};

const DEFAULT_POLL_INTERVAL: Duration = Duration::from_secs(1);

/// The core runner responsible for locking and running jobs
pub struct Runner<Context> {
    connection_pool: Pool<AsyncPgConnection>,
    queues: HashMap<String, Queue<Context>>,
    context: Context,
    shutdown_when_queue_empty: bool,
}

impl<Context: Clone + Send + Sync + 'static> Runner<Context> {
    pub fn new(connection_pool: Pool<AsyncPgConnection>, context: Context) -> Self {
        Self {
            connection_pool,
            queues: HashMap::new(),
            context,
            shutdown_when_queue_empty: false,
        }
    }

    /// Registers a new job type for this job runner.
    pub fn register_job_type<J: BackgroundJob<Context = Context>>(mut self) -> Self {
        let queue = self.queues.entry(J::QUEUE.into()).or_default();
        queue.job_registry.register::<J>();
        self
    }

    /// Adjusts the configuration of the [`DEFAULT_QUEUE`] queue.
    pub fn configure_default_queue<F>(self, f: F) -> Self
    where
        F: FnOnce(&mut Queue<Context>) -> &Queue<Context>,
    {
        self.configure_queue(DEFAULT_QUEUE, f)
    }

    /// Adjusts the configuration of a queue. If the queue does not exist,
    /// it will be created.
    pub fn configure_queue<F>(mut self, name: &str, f: F) -> Self
    where
        F: FnOnce(&mut Queue<Context>) -> &Queue<Context>,
    {
        f(self.queues.entry(name.into()).or_default());
        self
    }

    /// Sets the runner to shut down when the background job queue is empty.
    pub fn shutdown_when_queue_empty(mut self) -> Self {
        self.shutdown_when_queue_empty = true;
        self
    }

    /// Starts the background workers.
    ///
    /// This returns a `RunHandle` which can be used to wait for the workers to shutdown.
    pub fn start(&self) -> RunHandle {
        // Workers wait on this to be woken up when a new job is enqueued.
        let notify = Arc::new(Notify::new());

        // Listen for `NOTIFY` messages and wake up the workers. This is skipped
        // when shutting down on an empty queue (e.g. in tests), where workers
        // simply poll until the queue drains.
        if !self.shutdown_when_queue_empty {
            let pool = self.connection_pool.clone();
            let notify = notify.clone();
            tokio::spawn(async move {
                // `listen_for_new_jobs()` loops forever, so it should never
                // return. Catch a panic (or an unexpected return) and log it,
                // since the workers otherwise silently fall back to polling.
                if AssertUnwindSafe(listener::listen_for_new_jobs(pool, notify))
                    .catch_unwind()
                    .await
                    .is_err()
                {
                    error!("Background job listener panicked");
                } else {
                    warn!("Background job listener stopped unexpectedly");
                }
            });
        }

        let mut handles = Vec::new();
        for (queue_name, queue) in &self.queues {
            for i in 1..=queue.num_workers {
                let name = format!("background-worker-{queue_name}-{i}");
                info!(worker.name = %name, "Starting worker…");

                let worker = Worker {
                    connection_pool: self.connection_pool.clone(),
                    context: self.context.clone(),
                    job_registry: Arc::new(queue.job_registry.clone()),
                    shutdown_when_queue_empty: self.shutdown_when_queue_empty,
                    poll_interval: queue.poll_interval,
                    notify: notify.clone(),
                };

                let span = info_span!("worker", worker.name = %name);
                let handle = tokio::spawn(async move { worker.run().instrument(span).await });

                handles.push(handle);
            }
        }

        RunHandle { handles }
    }

    /// Checks if any jobs in the queue have failed.
    ///
    /// This function is intended for use in tests and will return an error if
    /// any jobs have failed.
    pub async fn check_for_failed_jobs(&self) -> anyhow::Result<()> {
        let mut conn = self.connection_pool.get().await?;

        let failed_jobs = storage::failed_job_count(&mut conn).await?;
        if failed_jobs == 0 {
            Ok(())
        } else {
            Err(anyhow!("{failed_jobs} jobs failed"))
        }
    }
}

pub struct RunHandle {
    handles: Vec<JoinHandle<()>>,
}

impl RunHandle {
    /// Waits for all background workers to shut down.
    pub async fn wait_for_shutdown(self) {
        join_all(self.handles).await.into_iter().for_each(|result| {
            if let Err(error) = result {
                warn!("Background worker task panicked: {error}");
            }
        });
    }
}

pub struct Queue<Context> {
    job_registry: JobRegistry<Context>,
    num_workers: usize,
    poll_interval: Duration,
}

impl<Context> Default for Queue<Context> {
    fn default() -> Self {
        Self {
            job_registry: JobRegistry::default(),
            num_workers: 1,
            poll_interval: DEFAULT_POLL_INTERVAL,
        }
    }
}

impl<Context> Queue<Context> {
    /// Sets the number of workers to spawn for this queue.
    pub fn num_workers(&mut self, num_workers: usize) -> &mut Self {
        self.num_workers = num_workers;
        self
    }

    /// Sets the interval after which each worker of this queue polls for new jobs.
    pub fn poll_interval(&mut self, poll_interval: Duration) -> &mut Self {
        self.poll_interval = poll_interval;
        self
    }
}
