// SPDX-License-Identifier: AGPL-3.0-or-later

use std::{
    io::IsTerminal,
    process::{ExitCode, Termination},
};

use anyhow::{Context, anyhow};
use fork::{Fork, daemon};
use futures::TryFutureExt;
use time::{
    UtcOffset,
    format_description::{self, BorrowedFormatItem},
};
use tokio::{net::UnixStream, process::Command, runtime};
use tracing::{debug, info, trace, warn};
use tracing_subscriber::{
    fmt::{
        format::{Compact, DefaultFields, Format},
        time::OffsetTime,
    },
    layer::SubscriberExt,
    util::SubscriberInitExt,
};

use crate::{
    args::{LogTarget, RuntimeDir, SubCommand},
    client::Client,
    server::{InitializeExit, Server, ServerArgs},
};
mod args;
mod client;
mod cmd;
mod executor;
mod report;
mod server;
mod stats;
mod transport;

fn main() -> MainResult {
    let args = args::Args::parse();

    // Determine whether we're going to detach, and how to construct the async runtime.
    // The server gets a multithreaded runtime, since it runs several long-lived,
    // independent tasks.
    let (detach, mut runtime) = if let args::SubCommand::Start { detach, .. } = args.cmd {
        (detach, {
            let mut builder = runtime::Builder::new_multi_thread();
            builder.enable_all();
            builder
        })
    } else {
        (false, {
            let mut builder = runtime::Builder::new_current_thread();
            builder.enable_all();
            builder
        })
    };

    // Main-in-practice, a.k.a. async-main
    // -----------------------------------
    // Since the server path may potentially fork, we cannot run any
    // tokio-related code until *after* the fork. So, here, we just define the
    // future to pass to the runtime, which we will instantiate either in this
    // process or in the forked child if we are daemonizing.
    let async_main = async {
        let tracing_subscriber = tracing_subscriber::registry().with(
            tracing_subscriber::filter::LevelFilter::from_level(args.verbosity),
        );

        if matches!(args.log, LogTarget::Syslog)
            || (!matches!(args.log, LogTarget::Stderr)
                && matches!(args.cmd, SubCommand::Start { detach: true, .. }))
        {
            match syslog_layer() {
                Ok(layer) => tracing_subscriber.with(layer).init(),
                Err(e) => {
                    eprintln!("Failed to initialize syslog. Falling back to stderr: {e:#}");
                    tracing_subscriber.with(stderr_layer()).init();
                }
            }
        } else {
            tracing_subscriber.with(stderr_layer()).init();
        }

        let _ = tracing::span!(tracing::Level::TRACE, "root").enter();
        trace!("Parsed CLI arguments: {args:?}");
        debug!(
            "Using runtime directory: {}",
            args.runtime_dir.dir().to_string_lossy(),
        );

        match args.cmd {
            args::SubCommand::Start {
                force, cmd_timeout, ..
            } => {
                let server_args = ServerArgs {
                    runtime_dir: &args.runtime_dir,
                    timeout: cmd_timeout,
                    force,
                };
                trace!("Server args: {server_args:?}");
                match Server::try_initialize(server_args).await {
                    Ok(server) => server.run().await,
                    Err(InitializeExit::AlreadyRunning(pid)) => {
                        if detach {
                            // Don't consider it an error for the server to already
                            // be running somewhere if we're detaching anyway.
                            Ok(())
                        } else {
                            // But if we want to run in the foreground, it is an
                            // error
                            Err(anyhow!("server already running (pid: {pid})"))
                        }
                    }
                    Err(InitializeExit::Error(e)) => Err(e),
                }
            }
            args::SubCommand::Run { mut cmd, timeout } => {
                if let Some(timeout) = timeout {
                    cmd.timeout(timeout);
                }
                let mut client = init_client(&args.runtime_dir, args.verbosity).await?;
                client.send_cmd(cmd).await
            }
            args::SubCommand::Flush { timeout } => {
                let mut client = init_client(&args.runtime_dir, args.verbosity).await?;
                client.send_flush(timeout).await
            }
            args::SubCommand::Stop { timeout } => {
                match Client::try_initialize(args.socketfile()).await {
                    Ok(mut client) => client.send_stop(timeout).await,
                    Err(e) => {
                        info!("Attempt to connect to server failed with: {e:#}");
                        warn!("Could not stop server, as it does not appear to be running");
                        Ok(())
                    }
                }
            }
        }
    };

    if detach {
        if let Ok(Fork::Child) = daemon(false, true) {
            runtime.build().unwrap().block_on(async_main).into()
        } else {
            Ok(()).into()
        }
    } else {
        runtime.build().unwrap().block_on(async_main).into()
    }
}

struct MainResult(Result<(), anyhow::Error>);
impl From<Result<(), anyhow::Error>> for MainResult {
    fn from(value: Result<(), anyhow::Error>) -> Self {
        Self(value)
    }
}
impl Termination for MainResult {
    fn report(self) -> std::process::ExitCode {
        match self.0 {
            Ok(v) => v.report(),
            Err(e) => {
                tracing::error!("{e:#}");
                ExitCode::FAILURE
            }
        }
    }
}

async fn init_client(
    runtime_dir: &RuntimeDir,
    verbosity: tracing::Level,
) -> Result<Client<UnixStream>, anyhow::Error> {
    Client::try_initialize(runtime_dir.socketfile())
        .or_else(async |_| {
            info!("No server appears to be running. Attempting to start one in the background");
            start_server(runtime_dir, verbosity).await?;

            let mut last_err = None;
            for _ in 0..20 {
                match Client::try_initialize(runtime_dir.socketfile()).await {
                    Ok(client) => return Ok(client),
                    Err(e) => {
                        last_err = Some(e);
                        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
                    }
                }
            }
            if let Some(e) = last_err {
                Err(e)
            } else {
                Err(anyhow!("could not start server"))
            }
        })
        .await
}

async fn start_server(
    runtime_dir: &RuntimeDir,
    verbosity: tracing::Level,
) -> Result<(), anyhow::Error> {
    trace!(
        ?runtime_dir,
        verbosity = verbosity.as_str(),
        "starting server in background"
    );
    let vees = if verbosity == tracing::Level::TRACE {
        Some("-vv")
    } else if verbosity == tracing::Level::DEBUG {
        Some("-v")
    } else {
        None
    };

    Command::new(std::env::current_exe().context("while determining server executable path")?)
        .arg("--runtime-dir")
        .arg(runtime_dir.dir())
        .args(["start", "--detach"])
        .args(vees)
        .spawn()
        .context("while spawning server process")?
        .wait()
        .await
        .context("while waiting for server to daemonize")?;

    Ok(())
}

type StdErrLayer<S> = tracing_subscriber::fmt::Layer<
    S,
    DefaultFields,
    Format<Compact, LogTime>,
    fn() -> std::io::Stderr,
>;

fn stderr_layer<S>() -> StdErrLayer<S> {
    tracing_subscriber::fmt::layer()
        .compact()
        .with_target(false)
        .with_timer(log_timer())
        .with_ansi(std::io::stderr().is_terminal())
        .with_writer(std::io::stderr)
}

type LogTime = OffsetTime<Vec<BorrowedFormatItem<'static>>>;

fn log_timer() -> LogTime {
    OffsetTime::new(
        UtcOffset::current_local_offset().unwrap_or(UtcOffset::UTC),
        format_description::parse("[year]-[month]-[day]T[hour]:[minute]:[second][offset_hour]")
            .expect("valid time format"),
    )
}

type SysLogFormatter<S> =
    tracing_subscriber::fmt::Layer<S, DefaultFields, Format<Compact, ()>, syslog_tracing::Syslog>;

fn syslog_layer<S>() -> Result<SysLogFormatter<S>, anyhow::Error> {
    let (options, facility) = Default::default();
    let writer = syslog_tracing::Syslog::new(c"procfarm", options, facility)
        .context("failed to create syslog writer")?;
    Ok(tracing_subscriber::fmt::layer()
        .compact()
        .with_target(false)
        .with_level(false)
        .without_time()
        .with_ansi(false)
        .with_writer(writer))
}
