prek 0.5.5

use std::cmp::max;
use std::ffi::{OsStr, OsString};
use std::path::Path;
use std::sync::LazyLock;

use futures_util::{StreamExt, TryStreamExt};
use prek_consts::env_vars::{EnvVars, EnvVarsRead};
use rustc_hash::FxHashMap;
use tracing::trace;

use crate::config::PassFilenames;
use crate::hook::Hook;
use crate::warn_user;

fn cpu_count() -> usize {
    std::thread::available_parallelism()
        .map(std::num::NonZero::get)
        .unwrap_or(1)
}

fn resolve_concurrency(env_vars: &impl EnvVarsRead, primary_env_var: &str) -> usize {
    if env_vars.is_set(EnvVars::PREK_NO_CONCURRENCY) {
        return 1;
    }

    let cpu = cpu_count();
    let Ok(value) = env_vars.var(primary_env_var) else {
        return cpu;
    };
    if let Ok(cap) = value.parse::<usize>() {
        return cap.max(1);
    }
    warn_user!(
        "Invalid value for {primary_env_var}: {value:?}. Expected a positive integer; using default ({cpu})"
    );

    cpu
}

pub(crate) static INTERNAL_CONCURRENCY: LazyLock<usize> = LazyLock::new(cpu_count);

pub(crate) static HOOK_CONCURRENCY: LazyLock<usize> =
    LazyLock::new(|| resolve_concurrency(&EnvVars, EnvVars::PREK_CONCURRENT_HOOKS));

pub(crate) static BATCH_CONCURRENCY: LazyLock<usize> =
    LazyLock::new(|| resolve_concurrency(&EnvVars, EnvVars::PREK_CONCURRENT_BATCHES));

/// Iterator that yields partitions of filenames that fit within the maximum command line length.
struct Partitions<'a> {
    filenames: &'a [&'a Path],
    current_index: usize,
    max_per_batch: usize,
    remaining_arg_length: usize,
}

/// We make a conservative guess for the size of a single pointer (64-bit) here
/// in order to support scenarios where a 32-bit binary is launching a 64-bit
/// binary.
const POINTER_SIZE_CONSERVATIVE: usize = 8;

/// POSIX requires that we leave 2048 bytes of space so that the child processes
/// can have room to set their own environment variables.
const ARG_HEADROOM: usize = 2048;

// Adapted from https://github.com/sharkdp/argmax
/// Required size for a single KEY=VAR environment variable string and the
/// corresponding pointer in envp**.
fn environment_variable_size<O: AsRef<OsStr>>(key: O, value: O) -> usize {
    POINTER_SIZE_CONSERVATIVE   // size for the pointer in envp**
        + key.as_ref().len()    // size for the variable name
        + 1                     // size for the '=' sign
        + value.as_ref().len()  // size for the value
        + 1 // terminating NULL
}

/// Required size to store a single ARG argument and the corresponding
/// pointer in argv**.
fn arg_size<O: AsRef<OsStr>>(arg: O) -> usize {
    POINTER_SIZE_CONSERVATIVE  // size for the pointer in argv**
        + arg.as_ref().len()   // size for argument string
        + 1 // terminating NULL
}

#[cfg(unix)]
static ARG_MAX: LazyLock<usize> = LazyLock::new(|| {
    let arg_max = unsafe { libc::sysconf(libc::_SC_ARG_MAX) };
    if arg_max <= 0 {
        1 << 12
    } else {
        usize::try_from(arg_max).expect("SC_ARG_MAX too large")
    }
});

#[cfg(unix)]
static PAGE_SIZE: LazyLock<usize> = LazyLock::new(|| {
    let page_size = unsafe { libc::sysconf(libc::_SC_PAGE_SIZE) };
    if page_size < 4096 {
        4096
    } else {
        usize::try_from(page_size).expect("SC_PAGE_SIZE too large")
    }
});

// https://www.in-ulm.de/~mascheck/various/argmax/
// https://cgit.git.savannah.gnu.org/cgit/findutils.git/tree/xargs/xargs.c
// https://github.com/rust-lang/rust/issues/40384
// https://github.com/uutils/findutils/blob/af48c151fe9b29cb7d25471b5388013ca15748ba/src/xargs/mod.rs#L177
// https://github.com/sharkdp/argmax
fn platform_max_cli_length() -> usize {
    cfg_select! {
        unix => {
            let mut arg_max = *ARG_MAX;
            // Assume arguments are counted with the granularity of a single page,
            // so allow a one page cushion to account for rounding up
            arg_max -= *PAGE_SIZE;
            // POSIX recommends an additional 2048 bytes of headroom
            arg_max -= ARG_HEADROOM;
            arg_max.clamp(1 << 12, 1 << 20)
        },
        windows => (1 << 15) - ARG_HEADROOM, // UNICODE_STRING max - headroom
        _ => 1 << 12,
    }
}

fn env_size(override_envs: &FxHashMap<String, String>) -> usize {
    std::env::vars_os()
        .map(|(key, value)| {
            if key
                .to_str()
                .map(|key| override_envs.contains_key(key))
                .unwrap_or(false)
            {
                // key is in override_envs; add it later.
                0
            } else {
                environment_variable_size(&key, &value)
            }
        })
        .sum::<usize>()
        + override_envs
            .iter()
            .map(|(key, value)| environment_variable_size(key, value))
            .sum::<usize>()
}

impl<'a> Partitions<'a> {
    fn split(
        hook: &'a Hook,
        entry: &[OsString],
        filenames: &'a [&'a Path],
        concurrency: usize,
    ) -> anyhow::Result<Self> {
        let max_per_batch = match hook.pass_filenames {
            PassFilenames::Limited(n) => n.get(),
            _ => max(4, filenames.len().div_ceil(concurrency)),
        };
        let mut arg_max = platform_max_cli_length();

        let cmd = Path::new(&entry[0]);
        if cfg!(windows)
            && cmd.extension().is_some_and(|ext| {
                ext.eq_ignore_ascii_case("cmd") || ext.eq_ignore_ascii_case("bat")
            })
        {
            // Reduce max length for batch files on Windows due to cmd.exe limitations.
            // 1024 is additionally subtracted to give headroom for further
            // expansion inside the batch file.
            arg_max = 8192 - 1024;
        } else if cfg!(unix) {
            // We have to share space with the environment variables
            arg_max -= env_size(&hook.env);
            // Account for the terminating NULL entry
            arg_max -= POINTER_SIZE_CONSERVATIVE;
        }

        let args_size = entry.iter().map(arg_size).sum::<usize>()
            + hook.args.iter().map(arg_size).sum::<usize>()
            + POINTER_SIZE_CONSERVATIVE; // terminating NULL

        if args_size >= arg_max {
            anyhow::bail!(
                "Command line length ({args_size} bytes) exceeds platform limit ({arg_max} bytes).
                \nhint: Shorten the hook `entry`/`args` or wrap the command in a script to reduce command-line length.",
            );
        }
        arg_max -= args_size;

        Ok(Self {
            filenames,
            current_index: 0,
            max_per_batch,
            remaining_arg_length: arg_max,
        })
    }
}

impl<'a> Iterator for Partitions<'a> {
    type Item = &'a [&'a Path];

    fn next(&mut self) -> Option<Self::Item> {
        // Handle empty filenames case
        if self.filenames.is_empty() && self.current_index == 0 {
            self.current_index = 1;
            return Some(&[]);
        }

        if self.current_index >= self.filenames.len() {
            return None;
        }

        let start_index = self.current_index;
        let mut remaining_length = self.remaining_arg_length;

        while self.current_index < self.filenames.len() {
            let filename = self.filenames[self.current_index];
            let length = arg_size(filename);

            if length > remaining_length || self.current_index - start_index >= self.max_per_batch {
                break;
            }

            remaining_length -= length;
            self.current_index += 1;
        }

        if self.current_index == start_index {
            // If we couldn't add even a single file to this batch, it means the file
            // is too long to fit in the command line by itself.
            let filename = self.filenames[self.current_index];
            let length = arg_size(filename);
            panic!(
                "Filename `{}` is too long ({length} bytes) to fit in command line (remaining {remaining_length} bytes).",
                filename.display(),
            );
        }

        Some(&self.filenames[start_index..self.current_index])
    }
}

#[derive(Debug, Default)]
pub(crate) struct HookRunOutput {
    exit_code: i32,
    output: Vec<u8>,
}

impl HookRunOutput {
    pub(crate) fn new(exit_code: i32, output: Vec<u8>) -> Self {
        Self { exit_code, output }
    }

    fn append(&mut self, output: Self) {
        self.exit_code |= output.exit_code;
        self.output.extend(output.output);
    }
}

pub(crate) async fn run_by_batch<T, F>(
    hook: &Hook,
    filenames: &[&Path],
    entry: &[OsString],
    run: F,
) -> anyhow::Result<(i32, Vec<u8>)>
where
    F: for<'a> AsyncFn(&'a [&'a Path]) -> anyhow::Result<T>,
    T: Into<HookRunOutput> + Send + 'static,
{
    let concurrency = if hook.require_serial {
        1
    } else {
        *BATCH_CONCURRENCY
    };

    // Split files into batches
    let partitions = Partitions::split(hook, entry, filenames, concurrency)?;
    trace!(
        total_files = filenames.len(),
        concurrency = concurrency,
        "Running {}",
        hook.id,
    );

    #[allow(clippy::redundant_closure)]
    let output = futures_util::stream::iter(partitions)
        .map(|batch| run(batch))
        .buffered(concurrency)
        .try_fold(
            HookRunOutput::default(),
            |mut combined, output| async move {
                combined.append(output.into());
                anyhow::Ok(combined)
            },
        )
        .await?;

    Ok((output.exit_code, output.output))
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::{Path, PathBuf};

    #[test]
    fn hook_run_output_append_bitwise_ors_exit_codes() {
        let mut output = HookRunOutput::new(1, Vec::new());

        output.append(HookRunOutput::new(3, Vec::new()));

        assert_eq!(output.exit_code, 3);
    }

    #[test]
    fn hook_run_output_append_preserves_output_order() {
        let mut output = HookRunOutput::new(0, b"first".to_vec());

        output.append(HookRunOutput::new(0, b"second".to_vec()));

        assert_eq!(output.output, b"firstsecond");
    }

    fn resolve_concurrency_from_map(values: &[(&str, &str)], primary_env_var: &str) -> usize {
        resolve_concurrency(&EnvVars::from_map(values), primary_env_var)
    }

    /// Helper to create a Partitions iterator for testing.
    /// This bypasses the Hook requirement by directly constructing the struct.
    fn create_test_partitions<'a>(
        filenames: &'a [&'a Path],
        remaining_arg_length: usize,
        max_per_batch: usize,
    ) -> Partitions<'a> {
        Partitions {
            filenames,
            current_index: 0,
            remaining_arg_length,
            max_per_batch,
        }
    }

    #[test]
    fn test_partitions_empty_filenames() {
        let filenames: Vec<&Path> = vec![];

        let mut partitions = create_test_partitions(&filenames, 4096, 10);

        // Should return empty slice once, then None
        let batch = partitions.next();
        assert!(batch.is_some());
        assert_eq!(batch.unwrap().len(), 0);

        let batch = partitions.next();
        assert!(batch.is_none());
    }

    #[test]
    #[should_panic(expected = "is too long")]
    fn test_partitions_long_filename_in_middle_panics() {
        let file1 = PathBuf::from("file1.txt");
        let long_name = "a".repeat(5000);
        let long_file = PathBuf::from(&long_name);
        let file3 = PathBuf::from("file3.txt");
        let filenames: Vec<&Path> = vec![&file1, &long_file, &file3];

        let mut partitions = create_test_partitions(&filenames, 1000, 10);

        // First batch should succeed with file1
        let batch1 = partitions.next();
        assert!(batch1.is_some());

        // Second batch should panic on the long filename
        // This ensures we don't silently skip file3
        partitions.next();
    }

    #[test]
    fn test_partitions_respects_max_per_batch() {
        // Create many small files
        let files: Vec<PathBuf> = (0..100)
            .map(|i| PathBuf::from(format!("f{i}.txt")))
            .collect();
        let file_refs: Vec<&Path> = files.iter().map(PathBuf::as_path).collect();

        let partitions = create_test_partitions(&file_refs, 100_000, 25);

        let all_batches: Vec<_> = partitions.map(<[&Path]>::len).collect();

        // Should have multiple batches due to max_per_batch
        assert!(all_batches.len() >= 4);

        // All files should have been processed
        let total_files: usize = all_batches.iter().sum();
        assert_eq!(total_files, 100);
    }

    #[test]
    fn test_resolve_concurrency_defaults_to_cpu() {
        assert_eq!(
            resolve_concurrency_from_map(&[], EnvVars::PREK_CONCURRENT_HOOKS),
            cpu_count()
        );
    }

    #[test]
    fn test_resolve_concurrency_uses_configured_limit() {
        assert_eq!(
            resolve_concurrency_from_map(
                &[(EnvVars::PREK_CONCURRENT_HOOKS, "4")],
                EnvVars::PREK_CONCURRENT_HOOKS,
            ),
            4
        );
    }

    #[test]
    fn test_resolve_concurrency_allows_above_cpu_limit() {
        assert_eq!(
            resolve_concurrency_from_map(
                &[(EnvVars::PREK_CONCURRENT_HOOKS, "9999")],
                EnvVars::PREK_CONCURRENT_HOOKS,
            ),
            9999
        );
    }

    #[test]
    fn test_resolve_concurrency_zero_floors_to_one() {
        assert_eq!(
            resolve_concurrency_from_map(
                &[(EnvVars::PREK_CONCURRENT_HOOKS, "0")],
                EnvVars::PREK_CONCURRENT_HOOKS,
            ),
            1
        );
    }

    #[test]
    fn test_resolve_concurrency_invalid_falls_back() {
        assert_eq!(
            resolve_concurrency_from_map(
                &[(EnvVars::PREK_CONCURRENT_HOOKS, "abc")],
                EnvVars::PREK_CONCURRENT_HOOKS,
            ),
            cpu_count()
        );
    }

    #[test]
    fn test_resolve_concurrency_no_concurrency_overrides_limits() {
        assert_eq!(
            resolve_concurrency_from_map(
                &[
                    (EnvVars::PREK_NO_CONCURRENCY, "1"),
                    (EnvVars::PREK_CONCURRENT_HOOKS, "8"),
                ],
                EnvVars::PREK_CONCURRENT_HOOKS,
            ),
            1
        );
    }

    #[test]
    fn test_partitions_respects_cli_length_limit() {
        // Create files that will exceed CLI length limit
        let files: Vec<PathBuf> = (0..10)
            .map(|i| PathBuf::from(format!("file{i}.txt")))
            .collect();
        let file_refs: Vec<&Path> = files.iter().map(PathBuf::as_path).collect();

        // Set a small max_cli_length to force multiple batches
        let partitions = create_test_partitions(&file_refs, 100, 100);

        let all_batches: Vec<_> = partitions.map(<[&Path]>::len).collect();

        // Should have multiple batches due to CLI length limit
        assert!(all_batches.len() > 1);

        // All files should have been processed
        let total_files: usize = all_batches.iter().sum();
        assert_eq!(total_files, 10);
    }
}