use futures::FutureExt;
use hydrant::config::Config;
use hydrant::control::{ApiBind, ApiBinds, DatabasePoisoned, Hydrant};
use miette::IntoDiagnostic;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::sync::Arc;
use std::time::Duration;
use tokio::signal::unix::{SignalKind, signal};
use tokio::sync::Notify;
use tracing::{error, info};

mod allocator;

const DEFAULT_API_PORT: u16 = 3000;
const DEFAULT_DEBUG_PORT: u16 = DEFAULT_API_PORT + 1;

fn default_api_binds() -> ApiBinds {
    ApiBinds::try_from_iter([
        SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), DEFAULT_API_PORT),
        SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), DEFAULT_API_PORT),
    ])
    .expect("two-element literal is non-empty")
}

const USAGE: &str = "\
usage: hydrant [--version | --help]

hydrant takes no other arguments, it reads its config from HYDRANT_* env vars.
see docs/configuration.md for the full list.
";

#[derive(Debug, PartialEq, Eq)]
enum Command {
    Run,
    Version,
    Help,
}

fn parse_args(mut args: impl Iterator<Item = std::ffi::OsString>) -> miette::Result<Command> {
    let Some(arg) = args.next() else {
        return Ok(Command::Run);
    };
    let command = match arg.to_str() {
        Some("--version" | "-V") => Command::Version,
        Some("--help" | "-h") => Command::Help,
        _ => miette::bail!(
            help = "config goes in HYDRANT_* env vars, see `hydrant --help`",
            "unexpected argument {}",
            arg.to_string_lossy()
        ),
    };
    if let Some(extra) = args.next() {
        miette::bail!("unexpected argument {}", extra.to_string_lossy());
    }
    Ok(command)
}

struct AppConfig {
    api_binds: Option<ApiBinds>,
    enable_debug: bool,
    debug_port: u16,
}

impl AppConfig {
    fn from_env() -> miette::Result<Self> {
        use hydrant::__cfg as cfg;
        let env = |key: &str| std::env::var(key).ok();
        let api_binds = ApiBinds::from_lookup(env, Some(default_api_binds()))?;
        let enable_debug = cfg!(env, "ENABLE_DEBUG", false);
        let debug_port_default = api_binds
            .iter()
            .flat_map(ApiBinds::iter)
            .find_map(ApiBind::tcp)
            .and_then(|addr| addr.port().checked_add(1))
            .unwrap_or(DEFAULT_DEBUG_PORT);
        let debug_port = cfg!(env, "DEBUG_PORT", debug_port_default);
        Ok(Self {
            api_binds,
            enable_debug,
            debug_port,
        })
    }
}

// long enough for a busy database to finish what it's writing, short enough that a
// supervisor's own stop timeout (systemd's is 90s) doesn't kill it first
const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(30);
const DRAIN_TIMEOUT: Duration = Duration::from_secs(2);

fn main() -> miette::Result<()> {
    // config lives in env vars, so a stray argument (a `--version` probe, a typo) would
    // otherwise boot a full default instance without anyone noticing
    match parse_args(std::env::args_os().skip(1))? {
        Command::Run => {}
        Command::Version => {
            println!("hydrant {}", env!("CARGO_PKG_VERSION"));
            return Ok(());
        }
        Command::Help => {
            print!("{USAGE}");
            return Ok(());
        }
    }

    rustls::crypto::aws_lc_rs::default_provider()
        .install_default()
        .ok();

    // logging starts before the config is read, so what reading it warns about gets
    // printed, and after `.env` is loaded, since RUST_LOG may be set there
    hydrant::config::env::load_dotenv();
    let directives = std::env::var("RUST_LOG")
        .ok()
        .filter(|v| !v.trim().is_empty());
    let env_filter = tracing_subscriber::EnvFilter::builder()
        .with_default_directive(tracing::Level::INFO.into())
        .parse_lossy(directives.as_deref().unwrap_or(hydrant::DEFAULT_LOG_FILTER));
    tracing_subscriber::fmt().with_env_filter(env_filter).init();

    let cfg = Config::from_lookup(|key| std::env::var(key).ok())?;
    let app = AppConfig::from_env()?;

    let stop = watch_stop_signals()?;
    let runtime = tokio::runtime::Runtime::new().into_diagnostic()?;
    let hydrant = runtime.block_on(Hydrant::new(cfg))?;

    let result = runtime.block_on(async {
        let api_fut = app
            .api_binds
            .map(|binds| hydrant.serve(binds).boxed())
            .unwrap_or_else(|| std::future::pending().boxed());
        let debug_fut = app
            .enable_debug
            .then(|| hydrant.serve_debug(app.debug_port).boxed())
            .unwrap_or_else(|| std::future::pending().boxed());
        let result = tokio::select! {
            r = hydrant.run()? => r,
            r = api_fut => r,
            r = debug_fut => r,
            _ = stop.notified() => Ok(()),
        };
        // the select dropped the listeners. reconnects go to the next hydrant
        hydrant.drain_streams(DRAIN_TIMEOUT).await;
        result
    });

    // dropping the runtime would wait on blocking tasks that only `Hydrant::shutdown` ends
    runtime.shutdown_background();
    info!("stopping, waiting up to {SHUTDOWN_TIMEOUT:?} for the database to close");
    let closed = hydrant.shutdown(SHUTDOWN_TIMEOUT);
    if let (Err(_), Err(e)) = (&result, &closed) {
        error!("{e}");
    }
    // a supervisor can tell this apart from other failures by the exit code
    if let Err(e) = &result
        && e.downcast_ref::<DatabasePoisoned>().is_some()
    {
        eprintln!("{e:?}");
        std::process::exit(10);
    }
    result?;
    closed?;
    info!("stopped");
    Ok(())
}

/// the first SIGINT or SIGTERM fires the returned notify, a second one exits on the spot.
/// it gets its own thread and runtime, since the main runtime is shut down while the
/// database closes and a signal then still has to reach someone.
fn watch_stop_signals() -> miette::Result<Arc<Notify>> {
    let stop = Arc::new(Notify::new());
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_io()
        .build()
        .into_diagnostic()?;
    // registered before returning, so a signal right after startup isn't the default kill
    let (mut int, mut term) = runtime
        .block_on(async {
            Ok::<_, std::io::Error>((
                signal(SignalKind::interrupt())?,
                signal(SignalKind::terminate())?,
            ))
        })
        .into_diagnostic()?;
    std::thread::Builder::new()
        .name("signals".into())
        .spawn({
            let stop = stop.clone();
            move || {
                runtime.block_on(async {
                    let mut next = async || {
                        tokio::select! {
                            _ = int.recv() => {}
                            _ = term.recv() => {}
                        }
                    };
                    next().await;
                    info!("got a stop signal, stopping (send another to exit right away)");
                    stop.notify_one();
                    next().await;
                    eprintln!("got a second stop signal, exiting without closing the database");
                    std::process::exit(1);
                })
            }
        })
        .into_diagnostic()?;
    Ok(stop)
}