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

use std::{
    fmt::Display,
    ops::Deref,
    path::{Path, PathBuf},
    sync::Arc,
    time::Duration,
};

use anyhow::{Context, anyhow};
use bitcode::{Decode, Encode};
use futures::{FutureExt, Stream, TryStreamExt, stream};
use tokio::{
    net::UnixListener,
    sync::{
        Notify,
        mpsc::{Sender, error::TrySendError},
    },
    task::JoinHandle,
};
use tracing::{debug, error, info, trace, warn};

use crate::{
    args::RuntimeDir,
    client::ClientMessage,
    cmd::Cmd,
    executor::Executor,
    report::ReportOptions,
    stats::Stats,
    transport::{Listener, Transport},
};

#[derive(Debug)]
pub struct ServerArgs<'a> {
    pub runtime_dir: &'a RuntimeDir,
    pub force: bool,
    pub timeout: Option<Duration>,
}
impl ServerArgs<'_> {
    fn socketfile(&self) -> &Path {
        self.runtime_dir.socketfile()
    }

    fn pidfile(&self) -> &Path {
        self.runtime_dir.pidfile()
    }
}

pub struct Server<T> {
    listener: T,
    executor: Sender<Cmd>,
    executor_handle: JoinHandle<Stats>,
    shutdown: Arc<Notify>,
    cmd_timeout: Option<Duration>,
    report_options: ReportOptions,
    // NB: we don't reference these after they are created, but it is important
    // we retain references to them, so they can be deleted on Drop when the
    // server exits.
    _socketfile: RuntimeFile,
    _pidfile: RuntimeFile,
}

pub enum InitializeExit {
    AlreadyRunning(i32),
    Error(anyhow::Error),
}
impl From<anyhow::Error> for InitializeExit {
    fn from(value: anyhow::Error) -> Self {
        Self::Error(value)
    }
}

impl Server<UnixListener> {
    pub async fn try_initialize(args: ServerArgs<'_>) -> Result<Self, InitializeExit> {
        info!("Initializing server");
        if args.force {
            info!("Forcibly stopping existing server, if running (force=true)");
            if let Some(pid) = pid_running(args.pidfile())? {
                terminate_pid(pid)
                    .map_err(|e| warn!("Failed to terminate existing server process: {e:#}"))
                    .ok();
            }
            std::fs::remove_file(args.pidfile())
                .map_err(|e| match e.kind() {
                    std::io::ErrorKind::NotFound => {}
                    _ => {
                        warn!("Atempt to remove existing pidfile failed: {e:#}")
                    }
                })
                .ok();
            std::fs::remove_file(args.socketfile())
                .map_err(|e| match e.kind() {
                    std::io::ErrorKind::NotFound => {}
                    _ => {
                        warn!("Atempt to remove existing socketfile failed: {e:#}")
                    }
                })
                .ok();
        }

        // Ensure the runtime directory exists.
        std::fs::create_dir_all(args.runtime_dir.dir())
            .context("Could not create runtime directory")?;

        // First, attempt to create & bind the socket. This will fail if another
        // server instance is running.
        info!("Binding to socket: {}", args.socketfile().to_string_lossy());
        let listener = match UnixListener::bind(args.socketfile()) {
            Ok(listener) => listener,
            Err(e) => match e.kind() {
                // This *probably* means the server is already running, but it
                // could also mean it got killed without cleaning up after itself.
                // So, check to be sure that the PID in the pidfile is actually
                // running before assuming we're running.
                std::io::ErrorKind::AddrInUse => {
                    debug!(
                        "Socket already exists, checking for running server instance: {}",
                        args.socketfile().to_string_lossy()
                    );
                    // There's a theoretical possibility we're running at ~exactly
                    // the same time as another initializing process, so give it
                    // that theoretical process a moment to write to the pidfile
                    // before checking.
                    tokio::time::sleep(Duration::from_millis(100)).await;
                    match pid_running(args.pidfile()) {
                        Ok(Some(pid)) => {
                            debug!("Another server instance is already running: {pid}");
                            return Err(InitializeExit::AlreadyRunning(pid));
                        }
                        Ok(None) => {
                            warn!(
                                "Socketfile exists, but no server instance running. \
                                 Attempting to take over the socket."
                            );
                            // This probably means another server didn't clean up,
                            // so go ahead and try to take over.
                            std::fs::remove_file(args.pidfile()).ok();
                            std::fs::remove_file(args.socketfile()).ok();
                            UnixListener::bind(args.socketfile())
                                .context("while binding to cleared out socket")?
                        }
                        Err(err) => {
                            return Err(err).with_context(
                                || "while checking previous server state due to: {e:#}",
                            )?;
                        }
                    }
                }
                _ => Err(e).context("when binding to the socket")?,
            },
        };

        debug!("Writing pidfile: {}", args.pidfile().to_string_lossy());
        // Write our PID to the pidfile.
        std::fs::write(args.pidfile(), std::process::id().to_string().as_bytes())
            .context("when writing pidfile")?;

        trace!("Spawning signal handler");
        let notify = Arc::new(Notify::new());
        tokio::spawn(handle_signals(notify.clone()));

        trace!("Spawning command executor");
        let (sender, receiver) = tokio::sync::mpsc::channel(1024);
        let executor = tokio::spawn(Executor::new(receiver).run());

        Ok(Self {
            listener,
            executor: sender,
            executor_handle: executor,
            shutdown: notify,
            cmd_timeout: args.timeout,
            report_options: ReportOptions::default(),
            _socketfile: args.socketfile().to_owned().into(),
            _pidfile: args.pidfile().to_owned().into(),
        })
    }
}

impl<T> Server<T>
where
    T: Listener,
{
    pub async fn run(mut self) -> Result<(), anyhow::Error> {
        let res = self.as_stream().try_for_each(async |_| Ok(())).await;

        self.finalize(Some(Duration::from_secs(10)))
            .await
            .map_err(|e| error!("{e:#}"))
            .ok();

        res
    }

    fn as_stream(&mut self) -> impl Stream<Item = Result<(), anyhow::Error>> {
        info!("Server running, awaiting messages");
        stream::try_unfold(self, async move |this| {
            // Wait for either a notification that we should shutdown or a new
            // socket connection.
            let socket = futures::select! {
                _ = this.shutdown.notified().fuse() => {
                    trace!("Shutting down server due to received signal");
                    return Ok(None);
                }
                socket = this.listener.accept().fuse() => {
                    socket.context("failed accepting new socket connection")?
                }
            };

            match this.handle_message(socket).await {
                Ok(op) => match op {
                    // Continue loop on Continue (duh)
                    ServerOperation::Continue => Ok(Some(((), this))),
                    // End loop on Stop
                    ServerOperation::Stop => Ok(None),
                },
                Err(e) => match e {
                    // Abort loop on fatal error
                    HandleMessageError::Fatal(e) => Err(e),
                    // Continue loop on non-fatal error
                    HandleMessageError::NonFatal(e) => {
                        warn!("Non-fatal server error: {e:#}");
                        Ok(Some(((), this)))
                    }
                },
            }
        })
    }

    async fn handle_message(
        &mut self,
        mut socket: impl Transport,
    ) -> Result<ServerOperation, HandleMessageError> {
        let message = socket
            .recv()
            .await
            .context("failed receiving message on socket")
            .map_err(HandleMessageError::non_fatal)?;

        match message {
            ServerMessage::Command(mut cmd) => {
                if let Some(timeout) = self.cmd_timeout {
                    cmd.default_timeout(timeout);
                }
                match self.executor.try_send(cmd) {
                    Ok(..) => socket
                        .send(&ClientMessage::Ack)
                        .await
                        .map_err(HandleMessageError::non_fatal)?,
                    Err(e) => match e {
                        TrySendError::Full(..) => socket
                            .send(&ClientMessage::Wait)
                            .await
                            .map_err(HandleMessageError::non_fatal)?,
                        TrySendError::Closed(_) => Err(HandleMessageError::fatal(anyhow!(
                            "server error: executor died"
                        )))?,
                    },
                };
                Ok(ServerOperation::Continue)
            }
            ServerMessage::Flush(timeout) => {
                trace!("Received Flush message");
                self.flush(timeout)
                    .await
                    .context("while flushing pending tasks")
                    .map_err(HandleMessageError::non_fatal)?;

                trace!("Sending ACK to client");
                socket
                    .send(&ClientMessage::Ack)
                    .await
                    .context("while sending flush ack")
                    .map_err(HandleMessageError::non_fatal)?;

                Ok(ServerOperation::Continue)
            }
            ServerMessage::Stop(timeout) => {
                self.pre_ack_shutdown(timeout)
                    .await
                    .map_err(HandleMessageError::non_fatal)?;

                socket
                    .send(&ClientMessage::Ack)
                    .await
                    .context("while sending stop ack")
                    .map_err(HandleMessageError::non_fatal)?;

                Ok(ServerOperation::Stop)
            }
        }
    }
}

impl<T> Server<T> {
    /// Allow any pending commands to complete, returning when there is nothing
    /// left in the queue.
    ///
    /// Optionally, apply a timeout, and return an error if not all processes
    /// have completed within the specified duration.
    async fn flush(&mut self, timeout: Option<Duration>) -> Result<(), anyhow::Error> {
        info!("Flushing any queued tasks");
        // Flushing is actually accomplished by swapping the existing Executor
        // with a new one, then waiting for the old one to finish any processing
        // it was doing.

        // First, we create a new Sender, so that we can take ownership of the
        // old one and drop it.
        let (mut sender, receiver) = tokio::sync::mpsc::channel(1024);

        // After this, `sender` is the old one, the one being flushed.
        std::mem::swap(&mut self.executor, &mut sender);

        // Create an empty, always pending future to swap in, so we can take
        // ownership of the one we're flushing.
        let mut executor_handle = tokio::spawn(futures::future::pending::<Stats>());

        // After this, `executor_handle` is the old one, the one being flushed.
        std::mem::swap(&mut self.executor_handle, &mut executor_handle);

        // Drop the Sender, await the executor, and get the stats to pass into the new one.
        let stats = Self::flush_executor(sender, executor_handle, timeout).await?;

        // Pass the stats and the new receiver into a new Executor, so it can pick
        // up where the old one left off.
        let mut executor_handle = tokio::spawn(Executor::with_stats(receiver, stats).run());

        // Swap out the Server-side channel and join handle for the old
        // executor with those of the new one.
        std::mem::swap(&mut self.executor_handle, &mut executor_handle);

        // After the swap above, `executor_handle` is the empty, alwaysp ending JoinHandle.
        // Drop it explicitly, to ensure nobody tries to use it for anything.
        drop(executor_handle);

        Ok(())
    }

    async fn finalize(self, timeout: Option<Duration>) -> Result<(), anyhow::Error> {
        let stats = Self::flush_executor(self.executor, self.executor_handle, timeout).await?;
        self.report_options.report_stats(&stats);
        Ok(())
    }

    async fn flush_executor(
        sender: Sender<Cmd>,
        handle: JoinHandle<Stats>,
        timeout: Option<Duration>,
    ) -> Result<Stats, anyhow::Error> {
        // Since we hold exclusive access to the Sender, we can drop it,
        // which means the recv() loop on the Executor side will consume anything
        // in its queue and then exit, which will end the Executor task.
        drop(sender);

        // When the Executor task finishes processing, it will exit, and the await
        // on its JoinHandle will complete.
        Ok(if let Some(timeout) = timeout {
            tokio::time::timeout(timeout, handle)
                .await
                .context("timeout while flushing pending tasks")??
        } else {
            handle.await?
        })
    }

    async fn pre_ack_shutdown(&mut self, timeout: Option<Duration>) -> Result<(), anyhow::Error> {
        info!("Flushing any in-progress commands");
        self.flush(timeout)
            .await
            .context("while performing pre-stop flush")?;
        Ok(())
    }
}

#[derive(Debug, PartialEq, Eq, Encode, Decode)]
pub enum ServerMessage {
    Command(Cmd),
    Flush(Option<Duration>),
    Stop(Option<Duration>),
}

enum ServerOperation {
    Continue,
    Stop,
}

enum HandleMessageError {
    Fatal(anyhow::Error),
    NonFatal(anyhow::Error),
}
impl HandleMessageError {
    fn fatal(e: impl Into<anyhow::Error>) -> Self {
        Self::Fatal(e.into())
    }

    fn non_fatal(e: impl Into<anyhow::Error>) -> Self {
        Self::NonFatal(e.into())
    }
}

/// Check if a process with the PID in the pidfile at `pidfile` is running.
///
/// `Ok(None)` means either there is no file at the path or the process
/// is not running.
///
/// `Ok(Some(pid))` means the process with the returned PID is running.
///
/// # Errors
/// - If an IOError is encountered when checking for existence of the pidfile
///   or reading the pidfile, it is returned
fn pid_running(pidfile: &Path) -> Result<Option<i32>, anyhow::Error> {
    use nix::sys::signal;
    use nix::unistd::Pid;

    trace!("Looking for pidfile: {}", pidfile.to_string_lossy());

    if std::fs::exists(pidfile).context("when checking for existence of pidfile")? {
        trace!("Reading contents of pidfile: {}", pidfile.to_string_lossy());
        let pid = std::fs::read(pidfile).context("when reading pidfile")?;
        let pid: i32 = String::from_utf8_lossy(&pid)
            .trim()
            .parse()
            .context("when parsing pid from pidfile")?;

        trace!(
            "Checking status of PID from pidfile: {pid}, {}",
            pidfile.to_string_lossy()
        );

        // send kill(0), no error means process is running
        Ok(signal::kill(Pid::from_raw(pid), None).ok().map(|_| pid))
    } else {
        trace!("No pidfile found at path: {}", pidfile.to_string_lossy());
        Ok(None)
    }
}

fn terminate_pid(pid: i32) -> Result<(), anyhow::Error> {
    use nix::sys::signal;
    use nix::unistd::Pid;

    signal::kill(Pid::from_raw(pid), signal::Signal::SIGTERM).context("when terminating pid")
}

async fn handle_signals(notify: Arc<Notify>) {
    use tokio::signal::unix::{SignalKind, signal};

    let signals = [
        (SignalKind::alarm(), "SIGALRM"),
        (SignalKind::hangup(), "SIGHUP"),
        (SignalKind::interrupt(), "SIGINT"),
        (SignalKind::pipe(), "SIGPIPE"),
        (SignalKind::quit(), "SIGQUIT"),
        (SignalKind::terminate(), "SIGTERM"),
        (SignalKind::user_defined1(), "SIGUSR1"),
        (SignalKind::user_defined2(), "SIGUSR2"),
    ];

    let signals = signals
        .into_iter()
        .map(|(s, name)| (signal(s).expect("failed to create signal handler"), name))
        .map(|(mut s, name)| {
            Box::pin(async move {
                s.recv().await;
                info!("Received {name}: shutting down");
            })
        });

    futures::future::select_all(signals).await;

    notify.notify_waiters();
}

/// A runtime file that gets automatically unlinked upon Drop.
struct RuntimeFile(PathBuf);

impl From<PathBuf> for RuntimeFile {
    fn from(value: PathBuf) -> Self {
        Self(value)
    }
}

impl Drop for RuntimeFile {
    fn drop(&mut self) {
        std::fs::remove_file(&self.0).ok();
    }
}

impl Deref for RuntimeFile {
    type Target = PathBuf;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl Display for RuntimeFile {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0.to_string_lossy())
    }
}
