// now: Nix-based distributed command runner
// Copyright (C) 2026 Eric Rodrigues Pires
//
// This program is free software: you can redistribute it and/or modify it under
// the terms of the GNU Affero General Public License as published by the Free
// Software Foundation, either version 3 of the License, or (at your option)
// any later version.
//
// This program is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
// FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for
// more details.
//
// You should have received a copy of the GNU Affero General Public License along
// with this program. If not, see <https://www.gnu.org/licenses/>.

use std::{
    ffi::{OsStr, OsString},
    io::Read as _,
    os::unix::ffi::{OsStrExt, OsStringExt},
    path::{Path, PathBuf},
    time::Duration,
};

use ahash::{HashSet, HashSetExt};
use async_trait::async_trait;
use color_eyre::eyre::{OptionExt, eyre};
use format_bytes::format_bytes;
use smol::{
    channel,
    io::{AsyncReadExt, AsyncWriteExt},
    lock::{
        RwLock, Semaphore,
        futures::{Acquire, Read, Write},
    },
    process::{Child, Command, Stdio},
};

use crate::{
    builder::{NixConfig, NowBuilder},
    job::JobError,
    utils::{get_random_string, wait_for_output, write_output_to_stderr},
    workflow::NowCheckout,
};

pub(crate) struct SshOptions {
    pub(crate) uri: String,
    pub(crate) control_path: PathBuf,
    pub(crate) identity: Option<String>,
    pub(crate) known_hosts: Option<PathBuf>,
    pub(crate) keep_alive: Option<String>,
}

pub(crate) struct RemoteBuilder {
    pub(crate) cancellation: channel::Sender<()>,
    pub(crate) receiver: channel::Receiver<()>,
    pub(crate) semaphore: Semaphore,
    pub(crate) lock: RwLock<()>,
    pub(crate) ssh_options: SshOptions,
    pub(crate) host_system: String,
    pub(crate) build_systems: HashSet<String>,
    pub(crate) system_features: HashSet<String>,
    pub(crate) required_features: HashSet<String>,
}

pub(crate) fn get_ssh_options(
    SshOptions {
        control_path,
        identity: ssh_identity,
        known_hosts,
        keep_alive,
        ..
    }: &SshOptions,
) -> impl Iterator<Item = OsString> {
    let mut options = Vec::new();

    // Identity
    if let Some(ssh_identity) = ssh_identity {
        options.extend_from_slice(&["-i".into(), ssh_identity.into()]);
    }

    // Control master
    let mut control_path_option: OsString = "ControlPath=".into();
    control_path_option.push(control_path);
    options.extend_from_slice(&[
        "-o".into(),
        "ControlMaster=auto".into(),
        "-o".into(),
        control_path_option,
        "-o".into(),
        "ControlPersist=300".into(),
        "-o".into(),
        "BatchMode=yes".into(),
    ]);

    // Keepalive
    if let Some(keep_alive) = keep_alive {
        let mut server_alive_interval: OsString = "ServerAliveInterval=".into();
        server_alive_interval.push(keep_alive);
        options.extend_from_slice(&[
            "-o".into(),
            server_alive_interval,
            "-o".into(),
            "ServerAliveCountMax=3".into(),
        ]);
    }

    // Known hosts
    if let Some(known_hosts) = known_hosts {
        let mut user_known_hosts_file_option: OsString = "UserKnownHostsFile=".into();
        user_known_hosts_file_option.push(known_hosts);
        options.extend_from_slice(&[
            "-o".into(),
            user_known_hosts_file_option,
            "-o".into(),
            "GlobalKnownHostsFile=/dev/null".into(),
            "-o".into(),
            "StrictHostKeyChecking=yes".into(),
        ]);
    }

    options.into_iter()
}

impl RemoteBuilder {
    pub(crate) async fn get_remote_builders(
        config: &NixConfig,
        builders: Option<String>,
        project_source: &Path,
        keep_alive: Option<Duration>,
    ) -> color_eyre::Result<Vec<Self>> {
        let builders = if let Some(builders) = builders {
            builders
        } else if let Some(file) = config.builders.value.strip_prefix('@') {
            let Ok(mut file) = smol::fs::File::open(file).await else {
                return Ok(vec![]);
            };
            let mut buf = Vec::new();
            file.read_to_end(&mut buf).await?;
            String::from_utf8(buf)?
        } else {
            config.builders.value.clone()
        };

        let mut vec = vec![];
        for builder in regex::Regex::new(r"[\n;]+")
            .expect("valid regex")
            .split(&builders)
        {
            let mut iter = builder.trim().split(' ');

            let Some(ssh_uri) = iter.next() else {
                continue;
            };
            let ssh_uri = match ssh_uri.split_once('?') {
                Some((left, _)) => left,
                None => ssh_uri,
            };
            let ssh_uri = if let Some(plain_uri) = ssh_uri.strip_prefix("ssh-ng://") {
                plain_uri.to_string()
            } else if let Some(plain_uri) = ssh_uri.strip_prefix("ssh://") {
                plain_uri.to_string()
            } else {
                continue;
            };

            let build_systems = if let Some(systems) = iter.next()
                && systems != "-"
            {
                systems
                    .split(',')
                    .map(std::string::ToString::to_string)
                    .collect()
            } else {
                [config.system.value.clone()].into_iter().collect()
            };

            let ssh_identity = iter.next().and_then(|identity| {
                if identity == "-" {
                    None
                } else {
                    Some(identity.to_string())
                }
            });

            let maximum_builds = if let Some(maximum_builds) = iter.next()
                && maximum_builds != "-"
            {
                str::parse(maximum_builds)?
            } else {
                1
            };

            let _speed_factor = iter.next();

            let system_features = if let Some(system_features) = iter.next()
                && system_features != "-"
            {
                system_features
                    .split(',')
                    .map(std::string::ToString::to_string)
                    .collect()
            } else {
                HashSet::new()
            };

            let required_features = if let Some(required_features) = iter.next()
                && required_features != "-"
            {
                required_features
                    .split(',')
                    .map(std::string::ToString::to_string)
                    .collect()
            } else {
                HashSet::new()
            };

            let remote_path = project_source.join(format!("remotes/ssh-{}", get_random_string(10)));
            smol::fs::create_dir_all(&remote_path).await?;

            let known_hosts = if let Some(known_hosts_key) = iter.next()
                && known_hosts_key != "-"
            {
                let known_hosts_path = remote_path.join("known_hosts");
                let mut decoder = base64::read::DecoderReader::new(
                    known_hosts_key.as_bytes(),
                    &base64::engine::general_purpose::STANDARD,
                );
                let mut buf = Vec::new();
                decoder.read_to_end(&mut buf)?;
                smol::fs::write(&known_hosts_path, &buf).await?;
                Some(known_hosts_path)
            } else {
                None
            };

            let (cancellation, receiver) = channel::bounded(1);

            let control_path = remote_path.join("socket");

            let ssh_options = SshOptions {
                uri: ssh_uri,
                control_path,
                identity: ssh_identity,
                known_hosts,
                keep_alive: keep_alive.map(|duration| duration.as_secs().to_string()),
            };

            // Get host system for remote
            let mut command = Command::new("ssh");
            for arg in get_ssh_options(&ssh_options) {
                command.arg(arg);
            }
            command
                .args([
                    &ssh_options.uri,
                    "nix",
                    "--extra-experimental-features",
                    "'nix-command flakes'",
                    "eval",
                    "--impure",
                    "--raw",
                    "--expr",
                    "builtins.currentSystem",
                ])
                .stdin(Stdio::null())
                .stdout(Stdio::piped())
                .stderr(Stdio::piped());

            let mut child = command.spawn()?;
            let output = wait_for_output(&mut child, None).await?;
            let host_system = if output.status.success() {
                String::from_utf8(output.stdout)?
            } else {
                write_output_to_stderr(&output)?;
                return Err(eyre!("Failed to connect to {}", ssh_options.uri));
            };

            vec.push(RemoteBuilder {
                cancellation,
                receiver,
                semaphore: Semaphore::new(maximum_builds),
                lock: RwLock::default(),
                ssh_options,
                host_system,
                build_systems,
                system_features,
                required_features,
            });
        }

        Ok(vec)
    }
}

#[async_trait(?Send)]
impl NowBuilder for RemoteBuilder {
    fn acquire(&self) -> (Acquire<'_>, Read<'_, ()>, &channel::Receiver<()>) {
        (self.semaphore.acquire(), self.lock.read(), &self.receiver)
    }

    fn acquire_exclusive(&self) -> (Write<'_, ()>, &channel::Receiver<()>) {
        (self.lock.write(), &self.receiver)
    }

    fn get_name(&self) -> String {
        self.ssh_options.uri.clone()
    }

    fn is_remote(&self) -> bool {
        true
    }

    async fn create_tempdir(
        &self,
        checkout: NowCheckout,
        cancellation: &channel::Receiver<()>,
    ) -> Result<Option<PathBuf>, JobError> {
        match checkout {
            NowCheckout::Default
            | NowCheckout::All
            | NowCheckout::None
            | NowCheckout::Clone
            | NowCheckout::CloneAll => {
                let mut tempdir_command = Command::new("ssh");
                for arg in get_ssh_options(&self.ssh_options) {
                    tempdir_command.arg(arg);
                }
                tempdir_command
                    .args([
                        &self.ssh_options.uri,
                        r#"sh -c 'mktemp -d "${TMPDIR:-/tmp}/now-XXXXXXXXXX"'"#,
                    ])
                    .stdin(Stdio::null())
                    .stdout(Stdio::piped())
                    .stderr(Stdio::piped());
                let mut child = tempdir_command.spawn().map_err(|error| eyre!(error))?;
                let output = wait_for_output(&mut child, Some(cancellation)).await?;
                if output.status.success() {
                    Ok(Some(PathBuf::from(OsStr::from_bytes(
                        output.stdout.trim_ascii(),
                    ))))
                } else {
                    write_output_to_stderr(&output)?;
                    Err(eyre!("Temporary directory creation failed").into())
                }
            }
        }
    }

    async fn checkout(
        &self,
        checkout: NowCheckout,
        tempdir: Option<PathBuf>,
        cancellation: &channel::Receiver<()>,
    ) -> Result<PathBuf, JobError> {
        let tempdir = tempdir.ok_or_eyre(eyre!("Missing temporary directory in remote"))?;

        match checkout {
            NowCheckout::Default | NowCheckout::Clone => {
                let mut ssh_command: OsString = "ssh ".into();
                for arg in get_ssh_options(&self.ssh_options) {
                    ssh_command.push(arg);
                    ssh_command.push(" ");
                }

                let mut command = Command::new("rsync");
                command
                    .arg("-e")
                    .arg(ssh_command)
                    .args(["-arz", "--files-from=-", "."])
                    .arg(OsString::from_vec(format_bytes!(
                        b"{}:{}",
                        self.ssh_options.uri.as_bytes(),
                        tempdir.as_os_str().as_encoded_bytes()
                    )))
                    .stdin(Stdio::piped())
                    .stdout(Stdio::piped())
                    .stderr(Stdio::piped());

                let mut child = command.spawn().map_err(|error| eyre!(error))?;
                let mut stdin = child.stdin.take().expect("stdin is piped");
                let stdin_future = async move {
                    let cwd = std::env::current_dir().map_err(|error| eyre!(error))?;
                    for dir_entry in
                        ignore::Walk::new(std::env::current_dir().map_err(|error| eyre!(error))?)
                            .flatten()
                    {
                        if dir_entry.file_type().is_some_and(|typ| typ.is_file()) {
                            let path = dir_entry.path();
                            stdin
                                .write_all(
                                    path.strip_prefix(&cwd)
                                        .unwrap_or(path)
                                        .as_os_str()
                                        .as_bytes(),
                                )
                                .await
                                .map_err(|error| eyre!(error))?;
                            stdin.write_all(b"\n").await.map_err(|error| eyre!(error))?;
                        }
                    }
                    Ok(stdin.flush().await.map_err(|error| eyre!(error))?)
                };
                let (_, output) = smol::future::try_zip(
                    stdin_future,
                    wait_for_output(&mut child, Some(cancellation)),
                )
                .await?;
                if output.status.success() {
                    Ok(tempdir)
                } else {
                    write_output_to_stderr(&output)?;
                    Err(eyre!("Checkout failed").into())
                }
            }
            NowCheckout::All | NowCheckout::CloneAll => {
                let mut ssh_command: OsString = "ssh ".into();

                for arg in get_ssh_options(&self.ssh_options) {
                    ssh_command.push(arg);
                    ssh_command.push(" ");
                }

                let mut command = Command::new("rsync");
                command
                    .arg("-e")
                    .arg(ssh_command)
                    .args(["-arz", "."])
                    .arg(OsString::from_vec(format_bytes!(
                        b"{}:{}",
                        self.ssh_options.uri.as_bytes(),
                        tempdir.as_os_str().as_encoded_bytes()
                    )))
                    .stdin(Stdio::null())
                    .stdout(Stdio::piped())
                    .stderr(Stdio::piped());

                let mut child = command.spawn().map_err(|error| eyre!(error))?;
                let output = wait_for_output(&mut child, Some(cancellation)).await?;
                if output.status.success() {
                    Ok(tempdir)
                } else {
                    write_output_to_stderr(&output)?;
                    Err(eyre!("Checkout failed").into())
                }
            }
            NowCheckout::None => Ok(tempdir),
        }
    }

    async fn copy_derivations(
        &self,
        job_name: &str,
        derivations: &[PathBuf],
        cancellation: &channel::Receiver<()>,
    ) -> Result<(), JobError> {
        let mut ssh_opts = OsString::new();
        for arg in get_ssh_options(&self.ssh_options) {
            ssh_opts.push(arg);
            ssh_opts.push(" ");
        }

        let mut command = Command::new("nix");
        command.args([
            "--extra-experimental-features",
            "nix-command flakes",
            "copy",
            "--to",
        ]);
        command
            .arg(format!("ssh-ng://{}", &self.ssh_options.uri))
            .args(derivations)
            .env("NIX_SSHOPTS", ssh_opts)
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());

        let mut child = command.spawn().map_err(|error| eyre!(error))?;
        let output = wait_for_output(&mut child, Some(cancellation)).await?;
        if !output.status.success() {
            write_output_to_stderr(&output)?;
            return Err(eyre!(
                "Failed to copy '{}' derivations to {}",
                job_name,
                self.ssh_options.uri
            )
            .into());
        }
        Ok(())
    }

    async fn realize_derivation(
        &self,
        derivation: &Path,
        cancellation: &channel::Receiver<()>,
    ) -> Result<PathBuf, JobError> {
        let mut full_command: OsString = OsString::new();

        if let Some(nix_config) = std::env::var_os("NIX_CONFIG") {
            full_command.push("NIX_CONFIG=");
            full_command.push(OsStr::from_bytes(
                shlex::bytes::try_quote(nix_config.as_encoded_bytes())
                    .map_err(|error| eyre!(error))?
                    .as_ref(),
            ));
            full_command.push(" ");
        }

        full_command.push("nix-store --realise ");
        full_command.push(derivation);

        let mut command = Command::new("ssh");
        for arg in get_ssh_options(&self.ssh_options) {
            command.arg(arg);
        }
        command
            .arg(&self.ssh_options.uri)
            .arg(full_command)
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());
        let mut child = command.spawn().map_err(|error| eyre!(error))?;
        let output = wait_for_output(&mut child, Some(cancellation)).await?;
        if output.status.success() {
            Ok(PathBuf::from(OsStr::from_bytes(output.stdout.trim_ascii())))
        } else {
            write_output_to_stderr(&output)?;
            Err(eyre!(
                "Failed to realize derivation '{}' in {}",
                derivation.to_string_lossy(),
                self.ssh_options.uri
            )
            .into())
        }
    }

    async fn download(
        &self,
        downloads: &[PathBuf],
        cancellation: &channel::Receiver<()>,
    ) -> Result<(), JobError> {
        let mut ssh_opts = OsString::new();
        for arg in get_ssh_options(&self.ssh_options) {
            ssh_opts.push(arg);
            ssh_opts.push(" ");
        }

        let mut command = Command::new("nix");
        command
            .args([
                "--extra-experimental-features",
                "nix-command flakes",
                "copy",
                "--to",
            ])
            .arg(format!("ssh-ng://{}", &self.ssh_options.uri))
            .args(downloads)
            .env("NIX_SSHOPTS", ssh_opts)
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());

        let mut child = command.spawn().map_err(|error| eyre!(error))?;
        let output = wait_for_output(&mut child, Some(cancellation)).await?;
        if !output.status.success() {
            write_output_to_stderr(&output)?;
            return Err(eyre!("Failed to copy uploads to {}", self.ssh_options.uri).into());
        }
        Ok(())
    }

    fn run_derivation(&self, cwdir: &Path, derivation: PathBuf) -> Result<Child, JobError> {
        let mut full_command: OsString = "cd ".into();
        full_command.push(cwdir);
        full_command.push(" ; ");

        full_command.push(derivation.join("bin/now-step"));

        let mut command = Command::new("ssh");
        for arg in get_ssh_options(&self.ssh_options) {
            command.arg(arg);
        }
        command
            .arg(&self.ssh_options.uri)
            .arg(full_command)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());

        Ok(command.spawn().map_err(|error| eyre!(error))?)
    }

    async fn fetch_derivation(
        &self,
        derivation: &Path,
        cancellation: &channel::Receiver<()>,
    ) -> Result<(), JobError> {
        let mut ssh_opts = OsString::new();
        for arg in get_ssh_options(&self.ssh_options) {
            ssh_opts.push(arg);
            ssh_opts.push(" ");
        }

        let mut command = Command::new("nix");
        command
            .args([
                "--extra-experimental-features",
                "nix-command flakes",
                "copy",
                "--no-check-sigs",
                "--from",
            ])
            .arg(format!("ssh-ng://{}", &self.ssh_options.uri))
            .arg(derivation)
            .env("NIX_SSHOPTS", ssh_opts)
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());

        let mut child = command.spawn().map_err(|error| eyre!(error))?;
        let output = wait_for_output(&mut child, Some(cancellation)).await?;
        if !output.status.success() {
            write_output_to_stderr(&output)?;
            return Err(eyre!(
                "Failed to copy '{}' derivation from {}",
                derivation.to_string_lossy(),
                self.ssh_options.uri
            )
            .into());
        }
        Ok(())
    }

    async fn undo_checkout(&self, checkout: NowCheckout, path: &Path) -> Result<(), JobError> {
        match checkout {
            NowCheckout::Default
            | NowCheckout::None
            | NowCheckout::Clone
            | NowCheckout::All
            | NowCheckout::CloneAll => {
                let mut rm_command: OsString = "rm -rf ".into();
                rm_command.push(path);

                let mut command = Command::new("ssh");
                for arg in get_ssh_options(&self.ssh_options) {
                    command.arg(arg);
                }
                command
                    .arg(&self.ssh_options.uri)
                    .arg(rm_command)
                    .stdin(Stdio::null())
                    .stdout(Stdio::piped())
                    .stderr(Stdio::piped());

                let mut child = command.spawn().map_err(|error| eyre!(error))?;
                let output = wait_for_output(&mut child, None).await?;
                if output.status.success() {
                    Ok(())
                } else {
                    write_output_to_stderr(&output)?;
                    Err(eyre!(
                        "Failed to remove checked out directory in {}",
                        self.ssh_options.uri
                    )
                    .into())
                }
            }
        }
    }
}