use std::sync::Arc;

use common_utils::{errors::CustomResult, id_type};
use diesel_models::enums::ProcessTrackerStatus;
use error_stack::{report, ResultExt};
use router_env::{
    instrument,
    tracing::{self, Instrument},
};
use time::Duration;
use tokio::sync::mpsc;

use super::{
    env::logger::{self, debug, error, warn},
    metrics,
};
use crate::{
    configs::settings::SchedulerSettings, errors, flow::SchedulerFlow,
    scheduler::SchedulerInterface, utils::*, SchedulerAppState, SchedulerSessionState,
};

#[instrument(skip_all)]
pub async fn start_producer<T, U, F>(
    state: &T,
    scheduler_settings: Arc<SchedulerSettings>,
    (tx, mut rx): (mpsc::Sender<()>, mpsc::Receiver<()>),
    app_state_to_session_state: F,
) -> CustomResult<(), errors::ProcessTrackerError>
where
    F: Fn(&T, &id_type::TenantId) -> CustomResult<U, errors::ProcessTrackerError>,
    T: SchedulerAppState,
    U: SchedulerSessionState,
{
    use std::time::Duration;

    use rand::distributions::{Distribution, Uniform};

    let mut rng = rand::thread_rng();

    // TODO: this can be removed once rand-0.9 is released
    // reference - https://github.com/rust-random/rand/issues/1326#issuecomment-1635331942
    #[allow(clippy::unnecessary_fallible_conversions)]
    let timeout = Uniform::try_from(0..=scheduler_settings.loop_interval)
        .change_context(errors::ProcessTrackerError::ConfigurationError)?;

    tokio::time::sleep(Duration::from_millis(timeout.sample(&mut rng))).await;

    let mut interval =
        tokio::time::interval(Duration::from_millis(scheduler_settings.loop_interval));

    let mut shutdown_interval = tokio::time::interval(Duration::from_millis(
        scheduler_settings.graceful_shutdown_interval,
    ));

    let signal = common_utils::signals::get_allowed_signals()
        .map_err(|error| {
            logger::error!("Signal Handler Error: {:?}", error);
            errors::ProcessTrackerError::ConfigurationError
        })
        .attach_printable("Failed while creating a signals handler")?;
    let handle = signal.handle();
    let task_handle =
        tokio::spawn(common_utils::signals::signal_handler(signal, tx).in_current_span());

    loop {
        match rx.try_recv() {
            Err(mpsc::error::TryRecvError::Empty) => {
                interval.tick().await;
                let tenants = state.get_tenants();
                for tenant in tenants {
                    let session_state = app_state_to_session_state(state, &tenant)?;
                    match run_producer_flow(&session_state, &scheduler_settings).await {
                        Ok(_) => (),
                        Err(error) => {
                            // Intentionally not propagating error to caller.
                            // Any errors that occur in the producer flow must be handled here only, as
                            // this is the topmost level function which is concerned with the producer flow.
                            error!(?error);
                        }
                    }
                }
            }
            Ok(()) | Err(mpsc::error::TryRecvError::Disconnected) => {
                logger::debug!("Awaiting shutdown!");
                rx.close();
                shutdown_interval.tick().await;
                logger::info!("Terminating producer");
                break;
            }
        }
    }
    handle.close();
    task_handle
        .await
        .change_context(errors::ProcessTrackerError::UnexpectedFlow)?;

    Ok(())
}

#[instrument(skip_all)]
pub async fn run_producer_flow<T>(
    state: &T,
    settings: &SchedulerSettings,
) -> CustomResult<(), errors::ProcessTrackerError>
where
    T: SchedulerSessionState,
{
    lock_acquire_release::<_, _, _>(state.get_db().as_scheduler(), settings, move || async {
        let tasks = fetch_producer_tasks(state.get_db().as_scheduler(), settings).await?;
        debug!("Producer count of tasks {}", tasks.len());

        // [#268]: Allow task based segregation of tasks
        let (cug_tasks, main_tasks): (Vec<storage::ProcessTracker>, Vec<storage::ProcessTracker>) =
            tasks.into_iter().partition(|task| {
                task.application_source == Some(diesel_models::enums::ApplicationSource::Cug)
            });

        divide_and_append_tasks(
            state.get_db().as_scheduler(),
            SchedulerFlow::Producer,
            main_tasks,
            settings,
            diesel_models::enums::ApplicationSource::Main,
        )
        .await?;
        divide_and_append_tasks(
            state.get_db().as_scheduler(),
            SchedulerFlow::Producer,
            cug_tasks,
            settings,
            diesel_models::enums::ApplicationSource::Cug,
        )
        .await?;

        Ok(())
    })
    .await?;

    Ok(())
}

#[instrument(skip_all)]
pub async fn fetch_producer_tasks(
    db: &dyn SchedulerInterface,
    conf: &SchedulerSettings,
) -> CustomResult<Vec<storage::ProcessTracker>, errors::ProcessTrackerError> {
    let upper = conf.producer.upper_fetch_limit;
    let lower = conf.producer.lower_fetch_limit;
    let now = common_utils::date_time::now();
    // Add these to validations
    let time_upper_limit = now.checked_add(Duration::seconds(upper)).ok_or_else(|| {
        report!(errors::ProcessTrackerError::ConfigurationError)
            .attach_printable("Error obtaining upper limit to fetch producer tasks")
    })?;
    let time_lower_limit = now.checked_sub(Duration::seconds(lower)).ok_or_else(|| {
        report!(errors::ProcessTrackerError::ConfigurationError)
            .attach_printable("Error obtaining lower limit to fetch producer tasks")
    })?;

    let mut new_tasks = db
        .find_processes_by_time_status(
            time_lower_limit,
            time_upper_limit,
            ProcessTrackerStatus::New,
            None,
        )
        .await
        .change_context(errors::ProcessTrackerError::ProcessFetchingFailed)?;
    let mut pending_tasks = db
        .find_processes_by_time_status(
            time_lower_limit,
            time_upper_limit,
            ProcessTrackerStatus::Pending,
            None,
        )
        .await
        .change_context(errors::ProcessTrackerError::ProcessFetchingFailed)?;

    if new_tasks.is_empty() {
        warn!("No new tasks found for producer to schedule");
    }
    if pending_tasks.is_empty() {
        warn!("No pending tasks found for producer to schedule");
    }

    new_tasks.append(&mut pending_tasks);

    // Safety: Assuming we won't deal with more than `u64::MAX` tasks at once
    #[allow(clippy::as_conversions)]
    metrics::TASKS_PICKED_COUNT.add(new_tasks.len() as u64, &[]);
    Ok(new_tasks)
}
