// 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::{
    env::temp_dir,
    ffi::{OsStr, OsString},
    num::NonZeroUsize,
    os::unix::ffi::OsStrExt,
    path::{Path, PathBuf},
    time::Duration,
};

use ahash::{HashMap, HashSet};
use async_trait::async_trait;
use color_eyre::eyre::{OptionExt, eyre};
use futures::{AsyncWriteExt, stream::FuturesUnordered};
use smol::{
    channel::{self, Receiver},
    lock::{
        RwLock, RwLockReadGuard, RwLockWriteGuard, Semaphore, SemaphoreGuard,
        futures::{Acquire, Read, Write},
    },
    process::{Child, Command, Stdio},
    stream::StreamExt,
};

use crate::{
    builder::{NixConfig, NowBuilder, remote::RemoteBuilder},
    environment::NowEnvironment,
    job::JobError,
    utils::{wait_for_output, write_output_to_stderr},
    workflow::{NowCheckout, NowJob, NowRunMode},
};

pub(crate) struct LocalBuilder {
    pub(crate) cancellation: channel::Sender<()>,
    pub(crate) receiver: channel::Receiver<()>,
    pub(crate) semaphore: Semaphore,
    pub(crate) lock: RwLock<()>,
    pub(crate) env: HashMap<OsString, OsString>,
    pub(crate) hostname: String,
    pub(crate) extra_platforms: Vec<String>,
    pub(crate) system: String,
    pub(crate) system_features: HashSet<String>,
    pub(crate) remote_builders: Vec<RemoteBuilder>,
    pub(crate) run_mode: NowRunMode,
    pub(crate) gcroot_dir: PathBuf,
}

pub(crate) struct BuilderGuard<'a> {
    pub(crate) guard: SemaphoreGuard<'a>,
    pub(crate) lock: RwLockReadGuard<'a, ()>,
    pub(crate) receiver: &'a Receiver<()>,
    pub(crate) builder: &'a dyn NowBuilder,
}

pub(crate) struct RunnerGuard<'a> {
    pub(crate) lock: RwLockWriteGuard<'a, ()>,
    pub(crate) receiver: &'a Receiver<()>,
    pub(crate) builder: &'a dyn NowBuilder,
}

impl LocalBuilder {
    pub(crate) async fn new(
        environment: &NowEnvironment,
        builders: Option<String>,
        run_mode: NowRunMode,
        cores: Option<NonZeroUsize>,
        keep_alive: Option<Duration>,
    ) -> color_eyre::Result<Self> {
        let mut command = Command::new("nix");
        command
            .args([
                "--extra-experimental-features",
                "nix-command flakes",
                "config",
                "show",
                "--json",
            ])
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());

        let mut child = command.spawn()?;
        let output = wait_for_output(&mut child, None).await?;
        let config: NixConfig = if output.status.success() {
            serde_json::from_slice(&output.stdout)?
        } else {
            write_output_to_stderr(&output)?;
            return Err(eyre!("Failed to fetch Nix config"));
        };

        let remote_builders = match run_mode {
            NowRunMode::All => {
                RemoteBuilder::get_remote_builders(
                    &config,
                    builders,
                    environment.nix_project_source.as_ref(),
                    keep_alive,
                )
                .await?
            }
            NowRunMode::LocalOnly => vec![],
            NowRunMode::RemoteOnly => {
                let remote_builders = RemoteBuilder::get_remote_builders(
                    &config,
                    builders,
                    environment.nix_project_source.as_ref(),
                    keep_alive,
                )
                .await?;
                if remote_builders.is_empty() {
                    return Err(eyre!("No remote builders available"));
                }
                remote_builders
            }
        };

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

        let cores = match cores {
            Some(cores) => cores.into(),
            None => smol::unblock(sysinfo::System::physical_core_count)
                .await
                .unwrap_or(1),
        };

        let hostname = smol::unblock(sysinfo::System::host_name)
            .await
            .unwrap_or_else(|| "localhost".into());

        let mut env = environment.local_env.clone();
        env.insert(
            "NOW_GCROOT_DIR".into(),
            environment.gcroot_dir.clone().into(),
        );

        Ok(Self {
            cancellation,
            receiver,
            semaphore: Semaphore::new(cores),
            lock: RwLock::default(),
            env,
            hostname,
            extra_platforms: config.extra_platforms.value,
            system: config.system.value,
            system_features: config.system_features.value.into_iter().collect(),
            remote_builders,
            run_mode,
            gcroot_dir: environment.gcroot_dir.clone(),
        })
    }

    pub(crate) fn cancel_builders(&self) {
        self.cancellation.close();
        for remote in &self.remote_builders {
            remote.cancellation.close();
        }
    }

    pub(crate) async fn get_builder(
        &self,
        job: &NowJob,
    ) -> color_eyre::Result<Option<BuilderGuard<'_>>> {
        let mut builders = vec![];

        match self.run_mode {
            NowRunMode::All | NowRunMode::LocalOnly => {
                if job.build_system == self.system
                    && job
                        .required_system_features
                        .iter()
                        .all(|feature| self.system_features.contains(feature))
                {
                    builders.push(self as &dyn NowBuilder);
                }
            }
            NowRunMode::RemoteOnly => {}
        }

        for builder in &self.remote_builders {
            if builder.build_systems.contains(&job.build_system)
                && builder
                    .required_features
                    .iter()
                    .all(|feature| job.required_system_features.contains(feature))
                && job
                    .required_system_features
                    .iter()
                    .all(|feature| builder.system_features.contains(feature))
            {
                builders.push(builder as &dyn NowBuilder);
            }
        }

        let mut builders_fut: FuturesUnordered<_> = builders
            .into_iter()
            .map(|builder| async move {
                let (semaphore, read, receiver) = builder.acquire();
                let lock = read.await;
                let guard = semaphore.await;
                BuilderGuard {
                    guard,
                    lock,
                    receiver,
                    builder,
                }
            })
            .collect();

        Ok(builders_fut.next().await)
    }

    pub(crate) fn has_runner(&self, job: &NowJob) -> bool {
        match self.run_mode {
            NowRunMode::All | NowRunMode::LocalOnly => {
                if (job.host_system == self.system
                    || self.extra_platforms.contains(&job.host_system))
                    && job
                        .required_system_features
                        .iter()
                        .all(|feature| self.system_features.contains(feature))
                {
                    return true;
                }
            }
            NowRunMode::RemoteOnly => {}
        }

        for builder in &self.remote_builders {
            if builder.host_system == job.host_system
                && builder
                    .required_features
                    .iter()
                    .all(|feature| job.required_system_features.contains(feature))
                && job
                    .required_system_features
                    .iter()
                    .all(|feature| builder.system_features.contains(feature))
            {
                return true;
            }
        }

        false
    }

    pub(crate) async fn get_runner(
        &self,
        job: &NowJob,
    ) -> color_eyre::Result<Option<RunnerGuard<'_>>> {
        let mut runners = vec![];

        match self.run_mode {
            NowRunMode::All | NowRunMode::LocalOnly => {
                if (job.host_system == self.system
                    || self.extra_platforms.contains(&job.host_system))
                    && job
                        .required_system_features
                        .iter()
                        .all(|feature| self.system_features.contains(feature))
                {
                    runners.push(self as &dyn NowBuilder);
                }
            }
            NowRunMode::RemoteOnly => {}
        }

        for builder in &self.remote_builders {
            if builder.host_system == job.host_system
                && builder
                    .required_features
                    .iter()
                    .all(|feature| job.required_system_features.contains(feature))
                && job
                    .required_system_features
                    .iter()
                    .all(|feature| builder.system_features.contains(feature))
            {
                runners.push(builder as &dyn NowBuilder);
            }
        }

        let mut runners_fut: FuturesUnordered<_> = runners
            .into_iter()
            .map(|builder| async move {
                let (write, receiver) = builder.acquire_exclusive();
                let lock = write.await;
                RunnerGuard {
                    lock,
                    receiver,
                    builder,
                }
            })
            .collect();

        Ok(runners_fut.next().await)
    }
}

#[async_trait(?Send)]
impl NowBuilder for LocalBuilder {
    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.hostname.clone()
    }

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

    async fn create_tempdir(
        &self,
        checkout: NowCheckout,
        cancellation: &channel::Receiver<()>,
    ) -> Result<Option<PathBuf>, JobError> {
        match checkout {
            NowCheckout::Default | NowCheckout::All => Ok(None),
            NowCheckout::None | NowCheckout::Clone | NowCheckout::CloneAll => {
                let mut tempdir_command = Command::new("sh");
                tempdir_command
                    .args(["-c", r#"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> {
        match checkout {
            NowCheckout::Default | NowCheckout::All => {
                debug_assert!(tempdir.is_none());
                Ok(std::env::current_dir().map_err(|error| eyre!(error))?)
            }
            NowCheckout::None => {
                let tempdir = tempdir.ok_or_eyre(eyre!("Missing temporary directory"))?;
                Ok(tempdir)
            }
            NowCheckout::Clone => {
                let tempdir = tempdir.ok_or_eyre(eyre!("Missing temporary directory"))?;
                let mut command = Command::new("rsync");
                command
                    .args(["-arz", "--files-from=-", "."])
                    .arg(&tempdir)
                    .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::CloneAll => {
                let tempdir = tempdir.ok_or_eyre(eyre!("Missing temporary directory"))?;

                let mut command = Command::new("rsync");
                command
                    .args(["-arz", "."])
                    .arg(&tempdir)
                    .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())
                }
            }
        }
    }

    async fn copy_derivations(
        &self,
        _job_name: &str,
        derivations: &[PathBuf],
        _cancellation: &channel::Receiver<()>,
    ) -> Result<(), JobError> {
        debug_assert!(derivations.iter().all(|path| path.exists()));
        Ok(())
    }

    async fn realize_derivation(
        &self,
        derivation: &Path,
        cancellation: &channel::Receiver<()>,
    ) -> Result<PathBuf, JobError> {
        let gcroot = self
            .gcroot_dir
            .join(derivation.file_name().expect("derivation has file name"));

        let mut command = Command::new("nix-store");
        command
            .arg("--add-root")
            .arg(&gcroot)
            .args(["--indirect", "--realise"])
            .arg(derivation)
            .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 '{}'",
                derivation.to_string_lossy(),
            )
            .into())
        }
    }

    async fn download(
        &self,
        downloads: &[PathBuf],
        _cancellation: &channel::Receiver<()>,
    ) -> Result<(), JobError> {
        debug_assert!(downloads.iter().all(|path| path.exists()));
        Ok(())
    }

    fn run_derivation(&self, cwdir: &Path, derivation: PathBuf) -> Result<Child, JobError> {
        let mut command = Command::new(derivation.join("bin/now-step"));
        command
            .current_dir(cwdir)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .env_clear()
            .envs(&self.env);
        Ok(command.spawn().map_err(|error| eyre!(error))?)
    }

    async fn fetch_derivation(
        &self,
        derivation: &Path,
        _cancellation: &channel::Receiver<()>,
    ) -> Result<(), JobError> {
        debug_assert!(derivation.exists());
        Ok(())
    }

    async fn undo_checkout(&self, checkout: NowCheckout, path: &Path) -> Result<(), JobError> {
        match checkout {
            NowCheckout::Default | NowCheckout::All => {
                debug_assert_eq!(
                    path.canonicalize().map_err(|error| eyre!(error))?,
                    std::env::current_dir()
                        .map_err(|error| eyre!(error))?
                        .canonicalize()
                        .map_err(|error| eyre!(error))?
                );
                Ok(())
            }
            NowCheckout::None | NowCheckout::Clone | NowCheckout::CloneAll => {
                assert_ne!(
                    path.canonicalize().map_err(|error| eyre!(error))?,
                    std::env::current_dir()
                        .map_err(|error| eyre!(error))?
                        .canonicalize()
                        .map_err(|error| eyre!(error))?
                );
                assert!(path.starts_with(temp_dir()));
                let mut command = Command::new("rm");
                command.arg("-rf").arg(path);

                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 locally checked out directory").into())
                }
            }
        }
    }
}