use anyhow::Result;
use chrono::Utc;
use tokio::sync::mpsc;
use tracing::info;

use rpki::rtr::admin::{
    AdminState, LogTailConfig, RuntimeConfigHandle, SourceReloadHandle, spawn_admin_config_server,
};
use rpki::rtr::bootstrap::{init_shared_cache, open_store, start_servers};
use rpki::rtr::config::{AppConfig, log_startup_config};
use rpki::rtr::report::{ReportConfiguration, ReportContext};
use rpki::rtr::runtime::spawn_refresh_task;
use rpki::rtr::server::{RtrService, RtrShutdownReason};
use rpki::slurm::admin::SlurmAdmin;

#[tokio::main]
async fn main() -> Result<()> {
    let config = AppConfig::from_env()?;
    init_tracing(config.timezone);
    log_startup_config(&config);

    let report_context = ReportContext::new(ReportConfiguration::new(
        config.source_refresh_interval.as_secs(),
        config.runtime_report_interval.as_secs(),
        config.report_history_limit,
        config.max_delta,
        config.prune_delta_by_snapshot_size,
        config.timezone,
        (
            config.timing.refresh,
            config.timing.retry,
            config.timing.expire,
        ),
    ));
    let store = open_store(&config)?;
    let shared_cache = init_shared_cache(&config, &store, &report_context)?;
    let runtime_config = RuntimeConfigHandle::new(config.runtime_config());

    let service = RtrService::with_config(shared_cache.clone(), config.service_config.clone());
    let notifier = service.notifier();
    let service_stats = service.stats();
    let shutdown_handle = service.shutdown_handle();
    let (source_reload_tx, source_reload_rx) = mpsc::channel(8);
    let source_reload = SourceReloadHandle::new(source_reload_tx);
    let (process_shutdown_tx, process_shutdown_rx) = mpsc::channel(1);

    let admin_task = config.admin_addr.map(|addr| {
        let slurm_admin = config.slurm_dir.as_ref().map(SlurmAdmin::new);
        let admin_state = AdminState::new(
            runtime_config.clone(),
            Some(source_reload.clone()),
            Some(shutdown_handle.clone()),
            Some(process_shutdown_tx.clone()),
            slurm_admin,
            LogTailConfig::from_env(),
        );
        spawn_admin_config_server(addr, config.admin_token.clone(), admin_state)
    });
    let running = start_servers(&config, &service);
    let refresh_task = spawn_refresh_task(
        &config,
        runtime_config,
        source_reload_rx,
        shared_cache.clone(),
        store.clone(),
        notifier,
        service_stats,
        report_context,
    );

    let stop_reason = wait_for_shutdown(process_shutdown_rx).await?;
    info!("stopping RTR service: reason={}", stop_reason.as_str());
    running.shutdown(stop_reason);
    running.wait().await;

    refresh_task.abort();
    let _ = refresh_task.await;
    if let Some(admin_task) = admin_task {
        admin_task.abort();
        let _ = admin_task.await;
    }

    info!("RTR service stopped");
    Ok(())
}

async fn wait_for_shutdown(
    mut process_shutdown_rx: mpsc::Receiver<RtrShutdownReason>,
) -> Result<RtrShutdownReason> {
    tokio::select! {
        signal = tokio::signal::ctrl_c() => {
            signal?;
            let reason = RtrShutdownReason::from_env();
            info!("shutdown signal received: reason={}", reason.as_str());
            Ok(reason)
        }
        reason = process_shutdown_rx.recv() => {
            let reason = reason.unwrap_or_else(RtrShutdownReason::from_env);
            info!("admin process lifecycle request received: reason={}", reason.as_str());
            Ok(reason)
        }
    }
}

fn init_tracing(timezone: chrono_tz::Tz) {
    let filter = tracing_subscriber::EnvFilter::try_from_default_env()
        .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn"));

    struct LocalTimer {
        timezone: chrono_tz::Tz,
    }

    impl tracing_subscriber::fmt::time::FormatTime for LocalTimer {
        fn format_time(
            &self,
            w: &mut tracing_subscriber::fmt::format::Writer<'_>,
        ) -> std::fmt::Result {
            let now = Utc::now().with_timezone(&self.timezone);
            write!(w, "{}", now.format("%Y-%m-%d %H:%M:%S%.3f %:z"))
        }
    }

    if let Err(err) = tracing_subscriber::fmt()
        .with_timer(LocalTimer { timezone })
        .with_env_filter(filter)
        .with_target(true)
        .with_thread_ids(true)
        .with_level(true)
        .try_init()
    {
        eprintln!("failed to initialize tracing subscriber: {err}");
    }
}