//! continuum: a replicated Automerge key-value store synchronized over Mycelium.
//!
//! main.rs only wires things together: arguments, state construction, background
//! tasks, bootstrap and the HTTP server.

use std::collections::HashMap;
use std::future::IntoFuture;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
use std::time::Duration;

use anyhow::{Context, Result};
use automerge::AutoCommit;
use clap::Parser;
use tokio::sync::{broadcast, mpsc, watch, RwLock};
use tracing::{debug, error, info, trace, warn};
use url::Url;

mod api;
mod document;
mod remake;
mod state;
mod swarm;
mod transport;

use api::auth::AuthState;
use document::replica::{Epoch, Replica};
use state::{AppState, NodeLifecycle, Timing};
use swarm::{bootstrap, peer, sync};
use transport::firewall::{Firewall, FirewallTransport};
use transport::mycelium::{self, MyceliumSetup, MyceliumSpawnConfig};
use transport::Transport;

#[derive(Parser, Debug)]
#[command(version, about = "continuum messaging system", long_about = None)]
struct Args {
    #[arg(long)]
    debug_api: bool,

    #[arg(long, value_name = "HOSTNAME:PORT", default_value = "127.0.0.1:3000")]
    bind_addr: SocketAddr,

    #[arg(
        long,
        value_name = "SCHEME://HOSTNAME:PORT",
        default_value = None,
        conflicts_with_all = &["mycelium_key_file", "mycelium_network_key_file", "mycelium_network_name"],
    )]
    mycelium_rest_url: Option<Url>,

    #[arg(long)]
    pub mycelium_key_file: Option<PathBuf>,

    #[arg(long)]
    pub mycelium_network_key_file: Option<PathBuf>,

    #[arg(long)]
    pub mycelium_network_name: Option<String>,

    #[arg(long)]
    pub mycelium_tcp_listen_port: Option<u16>,

    #[arg(
        long,
        use_value_delimiter = true,
        default_value = "tls://127.0.0.1:9651"
    )]
    pub mycelium_peers: Vec<String>,

    #[arg(long, value_name = "STRING,[STRING...]", value_delimiter = ',')]
    seeds: Option<Vec<String>>,

    #[arg(long, value_name = "PATH")]
    htpasswd: String,

    /// Pre-shared swarm key: only peers that know it may join (raw bytes, at least 16;
    /// e.g. `head -c 32 /dev/urandom`). Without it, anyone who knows the swarm ID can.
    #[arg(long, value_name = "PATH")]
    swarm_key_file: Option<PathBuf>,

    /// Keepalive interval; dormancy, resync and probe timings derive from it.
    #[arg(long, value_name = "SECONDS", default_value_t = 60)]
    keepalive_secs: u64,

    /// Test hook: start on a non-genesis epoch of this generation (with a fixed dummy
    /// id), to check that nodes on different epochs refuse to sync.
    #[arg(long, hide = true, requires = "debug_api", value_name = "GENERATION")]
    debug_epoch: Option<u64>,
}

#[tokio::main]
async fn main() -> Result<()> {
    console_subscriber::init();

    let args = Args::parse();

    // Read before anything starts, so a bad key file fails right away.
    let swarm_key = match &args.swarm_key_file {
        Some(path) => Some(Arc::new(swarm::auth::SwarmKey::from_file(path)?)),
        None => {
            warn!("No swarm key: any peer that learns the swarm ID can join");
            None
        }
    };

    // First thing: until a handler is installed, SIGINT and SIGTERM kill the process on
    // the spot, and startup (spawning Mycelium, bootstrap) can take minutes.
    let stop = watch_for_shutdown().context("Installing signal handlers")?;

    let mycelium_setup = {
        let mut spawn_config = MyceliumSpawnConfig {
            rest_url_opt: args.mycelium_rest_url,
            key_file_opt: args.mycelium_key_file,
            network_key_file: args.mycelium_network_key_file,
            peers: args.mycelium_peers,
            ..Default::default()
        };
        if let Some(net_name) = args.mycelium_network_name {
            spawn_config.network_name = net_name;
        }
        tokio::select! {
            setup = MyceliumSetup::bootstrap(spawn_config) => setup?,
            _ = stopped(stop.clone()) => return Ok(()),
        }
    };

    let mycelium_rest_url = mycelium_setup.rest_url.clone();

    info!(
        bind_addr = args.bind_addr.to_string(),
        "Starting replication service"
    );
    debug!(
        url = mycelium_rest_url.to_string(),
        "Using Mycelium REST API"
    );

    let seeds: Vec<String> = args.seeds.unwrap_or_default();
    debug!(seed_count = seeds.len(), "Loaded bootstrap configuration");

    let (sync_tx, sync_rx) = mpsc::unbounded_channel();
    let (change_tx, _) = watch::channel(0_u64);
    let (change_events, _) = broadcast::channel(256);

    let initial_epoch = match args.debug_epoch {
        Some(generation) => Epoch {
            generation,
            id: [0xee; 32],
        },
        None => Epoch::GENESIS,
    };
    let (epoch_tx, _) = watch::channel(initial_epoch);
    debug!(epoch = %initial_epoch, "Local epoch");

    let firewall = Arc::new(Firewall::default());

    let mycelium_transport = Arc::new(mycelium::MyceliumTransport::connect(mycelium_setup).await?);
    debug!(local_id = mycelium_transport.local_id(), "Transport ready");

    let state = AppState {
        doc: Arc::new(RwLock::new(Replica::new(initial_epoch, AutoCommit::new()))),
        epoch_tx,
        lifecycle: Arc::new(RwLock::new(NodeLifecycle::Joining {
            discovered_swarm_id: None,
        })),
        peer_states: Arc::new(RwLock::new(HashMap::new())),
        transport: Arc::new(FirewallTransport::new(
            mycelium_transport,
            Arc::clone(&firewall),
        )),
        firewall,
        sync_events: sync_tx,
        change_tx,
        initial_sync_done: Arc::new(AtomicBool::new(false)),
        change_events,
        timing: Timing::from_keepalive(Duration::from_secs(args.keepalive_secs)),
        remake: Arc::new(remake::RemakeState::default()),
        swarm_key,
        stale: Default::default(),
    };

    trace!("Application state initialized");

    let consumer_state = state.clone();
    tokio::spawn(async move {
        let inner_state = consumer_state.clone();
        let _ = consumer_state
            .transport
            .start_consumer(Box::new(move |sender, msg| {
                let state_clone = inner_state.clone();
                Box::pin(async move { peer::handle_message(&state_clone, sender, msg).await })
            }))
            .await;
    });
    trace!("Mycelium consumer task spawned");

    let sync_state = state.clone();
    tokio::spawn(async move {
        sync::run_sync_worker(sync_state, sync_rx).await;
    });
    trace!("Sync worker task spawned");

    let peer_worker_state = state.clone();
    tokio::spawn(async move {
        peer::run_peer_worker(peer_worker_state).await;
    });
    trace!("Peer worker task spawned");

    // Switches to committed epochs and catches up with newer ones. Starts before
    // bootstrapping: a node joining a compacted swarm can only finish its initial sync
    // after fetching the current epoch from a peer.
    tokio::spawn(remake::run_worker(state.clone()));
    trace!("Remake worker task spawned");

    tokio::select! {
        joined = bootstrap::bootstrap(&state, &seeds) => joined.context("Joining swarm")?,
        _ = stopped(stop.clone()) => {
            info!("Stopped while joining the swarm");
            return Ok(());
        }
    }

    // Start the change publisher after bootstrapping: streams only exist from then on,
    // and it avoids diffing the initial sync.
    let publisher_state = state.clone();
    tokio::spawn(async move { sync::run_change_publisher(publisher_state).await });

    info!(
        swarm_id = ?state.swarm_id().await,
        ready = state.is_ready().await,
        "Bootstrap completed"
    );

    let auth_state = AuthState::new(&args.htpasswd);
    let app = api::router(state.clone(), auth_state, args.debug_api);

    let listener = tokio::net::TcpListener::bind(args.bind_addr)
        .await
        .with_context(|| format!("Binding local API to {}", args.bind_addr))?;

    info!(address = %listener.local_addr()?, "HTTP API listening");

    let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();

    let server_builder = axum::serve(listener, app).with_graceful_shutdown(async move {
        let _ = shutdown_rx.await;
        info!("Shutdown signal received. Draining API connections...")
    });

    let server_fut = server_builder.into_future();
    tokio::pin!(server_fut);

    tokio::select! {
        exit_status = state.transport.wait_for_exit() => {
            error!("Mycelium runtime process exited. Status: {:?}", exit_status);
            let _ = shutdown_tx.send(());
            let _ = server_fut.as_mut().await;
        }

        server_result = server_fut.as_mut() => {
            if let Err(e) = server_result {
                error!("Axum HTTP server encountered a runtime error: {:?}", e);
            } else {
                info!("Axum HTTP engine shutdown cleanly.");
            }
            std::mem::drop(state);
        }

        _ = stopped(stop.clone()) => {
            let _ = shutdown_tx.send(());
            let _ = server_fut.as_mut().await;
        }
    }

    info!("HTTP API server stopped");
    Ok(())
}

/// Turns SIGINT and SIGTERM into a shutdown request. A second signal exits at once (the
/// supervised Mycelium still gets its parent-death signal then).
fn watch_for_shutdown() -> Result<watch::Receiver<bool>> {
    use tokio::signal::unix::{signal, SignalKind};

    let mut interrupt = signal(SignalKind::interrupt())?;
    let mut terminate = signal(SignalKind::terminate())?;
    let (stop_tx, stop_rx) = watch::channel(false);

    tokio::spawn(async move {
        let name = tokio::select! {
            _ = interrupt.recv() => "SIGINT",
            _ = terminate.recv() => "SIGTERM",
        };
        info!(
            signal = name,
            "Shutting down (send it again to exit immediately)"
        );
        let _ = stop_tx.send(true);

        tokio::select! {
            _ = interrupt.recv() => {}
            _ = terminate.recv() => {}
        }
        error!("Second signal; exiting immediately");
        std::process::exit(130);
    });

    Ok(stop_rx)
}

/// Resolves once a shutdown was requested.
async fn stopped(mut stop: watch::Receiver<bool>) {
    let _ = stop.wait_for(|stop| *stop).await;
}

// === Footnotes ===============================================================
// [concurrent]: these comments are in place so we remember what to do if we go from async to concurrent.