//! Functions to collect environment variables that are used during the build process.
use std::collections::HashMap;
use std::collections::HashSet;
use std::path::{Path, PathBuf};

use rattler_build_jinja::Variable;
use rattler_build_script::{EnvironmentIsolation, RuntimeEnv};
use rattler_build_types::NormalizedKey;
use rattler_conda_types::{Platform, RepoDataRecord};
use std::collections::BTreeMap;

use crate::android;
use crate::consts;
use crate::ios;
use crate::linux;
use crate::macos;
use crate::metadata::Output;
use crate::windows;

macro_rules! insert {
    ($map:expr, $key:expr, $value:expr) => {
        $map.insert($key.to_string(), Some($value.to_string()));
    };
}

fn get_stdlib_dir(prefix: &Path, platform: Platform, py_ver: &str) -> PathBuf {
    if platform.is_windows() {
        prefix.join("Lib")
    } else {
        let lib_dir = prefix.join("lib");
        lib_dir.join(format!("python{}", py_ver))
    }
}

fn get_sitepackages_dir(
    prefix: &Path,
    platform: Platform,
    py_ver: &str,
    python_site_packages_path: Option<&str>,
) -> PathBuf {
    python_site_packages_path.map_or_else(
        || get_stdlib_dir(prefix, platform, py_ver).join("site-packages"),
        |path| prefix.join(path),
    )
}

/// Returns a map of environment variables for Python that are used in the build process.
///
/// Variables:
/// - `PYTHON`: path to Python executable
/// - `PY3K`: 1 if Python 3, 0 if Python 2
/// - `PY_VER`: Python version (major.minor), e.g. 3.8
/// - `STDLIB_DIR`: Python standard library directory
/// - `SP_DIR`: Python site-packages directory
/// - `NPY_VER`: NumPy version (major.minor), e.g. 1.19
/// - `NPY_DISTUTILS_APPEND_FLAGS`: 1 (https://github.com/conda/conda-build/pull/3015)
pub fn python_vars(output: &Output) -> HashMap<String, Option<String>> {
    let mut result = HashMap::new();

    if output.host_platform().platform.is_windows() {
        let python = output.prefix().join("python.exe");
        insert!(result, "PYTHON", python.to_string_lossy());
    } else {
        let python = output.prefix().join("bin/python");
        insert!(result, "PYTHON", python.to_string_lossy());
    }

    // find python in the host dependencies
    let python_record = output
        .find_resolved_package("python")
        .map(|(record, _)| record);
    let python_version = output
        .variant()
        .get(&"python".into())
        .map(|s| s.to_string())
        .or_else(|| {
            // Use the resolved python version even if it's a transitive dependency
            // (e.g. pulled in by pip in noarch python packages)
            python_record.map(|record| record.package_record.version.to_string())
        });

    if let Some(py_ver) = python_version {
        let py_ver: Vec<_> = py_ver.split('.').take(2).collect();
        let py_ver_str = py_ver.join(".");
        let stdlib_dir = get_stdlib_dir(
            output.prefix(),
            output.host_platform().platform,
            &py_ver_str,
        );
        let site_packages_dir = get_sitepackages_dir(
            output.prefix(),
            output.host_platform().platform,
            &py_ver_str,
            python_record
                .and_then(|record| record.package_record.python_site_packages_path.as_deref()),
        );
        let py3k = if py_ver[0] == "3" { "1" } else { "0" };
        insert!(result, "PY3K", py3k);
        insert!(result, "PY_VER", py_ver_str);
        insert!(result, "STDLIB_DIR", stdlib_dir.to_string_lossy());
        insert!(result, "SP_DIR", site_packages_dir.to_string_lossy());
    }

    if let Some(npy_version) = output.variant().get(&"numpy".into()) {
        let npy_ver = npy_version.to_string();
        let npy_ver: Vec<_> = npy_ver.split('.').take(2).collect();
        let npy_ver = npy_ver.join(".");
        insert!(result, "NPY_VER", npy_ver);
        insert!(result, "NPY_DISTUTILS_APPEND_FLAGS", "1");
    }

    result
}

/// Returns Python-related environment variables derived from solved package records.
///
/// This is used in test environments where we have the solve result but no `Output`.
/// Extracts python and numpy versions from the resolved records.
pub fn python_vars_from_records(
    records: &[RepoDataRecord],
    prefix: &Path,
    platform: Platform,
) -> HashMap<String, Option<String>> {
    let mut result = HashMap::new();

    if platform.is_windows() {
        let python = prefix.join("python.exe");
        insert!(result, "PYTHON", python.to_string_lossy());
    } else {
        let python = prefix.join("bin/python");
        insert!(result, "PYTHON", python.to_string_lossy());
    }

    let python_record = records
        .iter()
        .find(|r| r.package_record.name.as_normalized() == "python");

    if let Some(python_record) = python_record {
        let py_ver = python_record.package_record.version.to_string();
        let py_ver: Vec<_> = py_ver.split('.').take(2).collect();
        let py_ver_str = py_ver.join(".");
        let stdlib_dir = get_stdlib_dir(prefix, platform, &py_ver_str);
        let site_packages_dir = get_sitepackages_dir(
            prefix,
            platform,
            &py_ver_str,
            python_record
                .package_record
                .python_site_packages_path
                .as_deref(),
        );
        let py3k = if py_ver[0] == "3" { "1" } else { "0" };
        insert!(result, "PY3K", py3k);
        insert!(result, "PY_VER", py_ver_str);
        insert!(result, "STDLIB_DIR", stdlib_dir.to_string_lossy());
        insert!(result, "SP_DIR", site_packages_dir.to_string_lossy());
    }

    let numpy_version = records
        .iter()
        .find(|r| r.package_record.name.as_normalized() == "numpy")
        .map(|r| r.package_record.version.to_string());

    if let Some(npy_ver) = numpy_version {
        let npy_ver: Vec<_> = npy_ver.split('.').take(2).collect();
        let npy_ver = npy_ver.join(".");
        insert!(result, "NPY_VER", npy_ver);
        insert!(result, "NPY_DISTUTILS_APPEND_FLAGS", "1");
    }

    result
}

/// Returns a map of environment variables for R that are used in the build process.
///
/// Variables:
/// - R_VER: R version (major.minor), e.g. 4.0
/// - R: Path to R executable
/// - R_USER: Path to R user directory
///
pub fn r_vars(output: &Output) -> HashMap<String, Option<String>> {
    let mut result = HashMap::new();

    if let Some(r_ver) = output.variant().get(&"r-base".into()) {
        insert!(result, "R_VER", r_ver);

        let r_bin = if output.host_platform().platform.is_windows() {
            output.prefix().join("Scripts/R.exe")
        } else {
            output.prefix().join("bin/R")
        };

        let r_user = output.prefix().join("Libs/R");

        insert!(result, "R", r_bin.to_string_lossy());
        insert!(result, "R_USER", r_user.to_string_lossy());
    }

    result
}

/// Returns a map of environment variables for all supported languages (Python, R).
pub fn language_vars(output: &Output) -> HashMap<String, Option<String>> {
    let mut result = HashMap::new();

    result.extend(python_vars(output));
    result.extend(r_vars(output));

    result
}

/// Returns a map of environment variables that are used in the build process.
/// Also adds platform-specific variables.
///
/// Variables:
/// - CPU_COUNT: Number of CPUs
/// - SHLIB_EXT: Shared library extension for platform (e.g. Linux -> .so, Windows -> .dll, macOS -> .dylib)
/// - LANG: Normalized to C.UTF-8 in strict mode, forwarded from host otherwise
/// - LC_ALL: Normalized to C.UTF-8 in strict mode, forwarded from host otherwise
///
/// Forwards the following environment variables:
/// - PATH: Path where executables are found
pub fn os_vars(
    prefix: &Path,
    target_platform: &Platform,
    host_platform: &Platform,
    build_platform: &Platform,
    env_isolation: EnvironmentIsolation,
    work_dir: &Path,
    runtime: &RuntimeEnv,
) -> HashMap<String, Option<String>> {
    let mut vars = HashMap::new();

    // For `noarch` outputs, fall back to the host platform so Windows target
    // vars (e.g. `LIBRARY_PREFIX`) are still emitted. See issue #2475.
    let os_platform = if *target_platform == Platform::NoArch {
        host_platform
    } else {
        target_platform
    };

    let path_var = if os_platform.is_windows() {
        "Path"
    } else {
        "PATH"
    };

    insert!(
        vars,
        "CPU_COUNT",
        runtime
            .var("CPU_COUNT")
            .map(str::to_owned)
            .unwrap_or_else(|| num_cpus::get().to_string())
    );

    match env_isolation {
        EnvironmentIsolation::Strict => {
            // Normalize locale for reproducible builds
            insert!(vars, "LANG", "C.UTF-8");
            insert!(vars, "LC_ALL", "C.UTF-8");
        }
        EnvironmentIsolation::CondaBuild | EnvironmentIsolation::None => {
            // Forward locale from host (conda-build behavior)
            vars.insert("LANG".to_string(), runtime.var("LANG").map(str::to_owned));
            vars.insert(
                "LC_ALL".to_string(),
                runtime.var("LC_ALL").map(str::to_owned),
            );
            vars.insert(
                "MAKEFLAGS".to_string(),
                runtime.var("MAKEFLAGS").map(str::to_owned),
            );
        }
    }

    insert!(
        vars,
        "SHLIB_EXT",
        rattler_build_types::shlib_ext(os_platform)
    );
    vars.insert(
        path_var.to_string(),
        runtime.var(path_var).map(str::to_owned),
    );

    if os_platform.is_windows() {
        vars.extend(windows::env::default_env_vars_target(prefix, runtime));
    } else if os_platform.is_osx() {
        vars.extend(macos::env::default_env_vars_target(
            prefix,
            os_platform,
            runtime,
        ));
    } else if os_platform.is_ios() {
        vars.extend(ios::env::default_env_vars_target(
            prefix,
            os_platform,
            runtime,
        ));
    } else if os_platform.is_android() {
        vars.extend(android::env::default_env_vars_target(
            prefix,
            os_platform,
            runtime,
        ));
    } else if os_platform.is_linux() {
        vars.extend(linux::env::default_env_vars_target(
            prefix,
            os_platform,
            env_isolation,
            runtime,
        ));
    }
    if build_platform.is_windows() {
        vars.extend(windows::env::default_env_vars_build(
            build_platform,
            runtime,
        ));
    } else if build_platform.is_osx() {
        vars.extend(macos::env::default_env_vars_build(build_platform));
    } else if build_platform.is_linux() {
        vars.extend(linux::env::default_env_vars_build(build_platform));
    }

    if build_platform.is_windows() {
        match env_isolation {
            EnvironmentIsolation::Strict => {
                insert!(vars, "USERPROFILE", work_dir.to_string_lossy());
            }
            EnvironmentIsolation::CondaBuild => {
                vars.insert(
                    "USERPROFILE".to_string(),
                    runtime.var("USERPROFILE").map(str::to_owned),
                );
            }
            EnvironmentIsolation::None => {}
        }
    } else {
        match env_isolation {
            EnvironmentIsolation::Strict => {
                insert!(vars, "HOME", work_dir.to_string_lossy());
                insert!(vars, "USER", "rattler");
                insert!(vars, "SHELL", "/bin/bash");
                insert!(vars, "EDITOR", "/bin/false");
                insert!(vars, "TERM", "xterm-256color");
            }
            EnvironmentIsolation::CondaBuild => {
                vars.insert(
                    "HOME".to_string(),
                    Some(runtime.var("HOME").unwrap_or("UNKNOWN").to_owned()),
                );
            }
            EnvironmentIsolation::None => {}
        }
    }

    vars
}

/// Set environment variables that help to force color output.
fn force_color_vars() -> HashMap<String, Option<String>> {
    let mut vars = HashMap::new();

    insert!(vars, "CLICOLOR_FORCE", "1");
    insert!(vars, "FORCE_COLOR", "1");
    insert!(vars, "AM_COLOR_TESTS", "always");
    insert!(vars, "MAKE_TERMOUT", "1");
    insert!(vars, "CMAKE_COLOR_DIAGNOSTICS", "ON");

    insert!(
        vars,
        "GCC_COLORS",
        "error=01;31:warning=01;35:note=01;36:caret=01;32:locus=01:quote=01"
    );

    vars
}

/// Return all variables that should be set during the build process, including
/// operating system specific environment variables.
pub fn vars(output: &Output, build_state: &str) -> HashMap<String, Option<String>> {
    let mut vars = HashMap::new();

    insert!(vars, "PYTHONNOUSERSITE", "1");

    if let Some((_, host_arch)) = output.host_platform().platform.to_string().rsplit_once('-') {
        insert!(vars, "ARCH", host_arch);
    }

    let directories = output.build_configuration.directories.exec_view();
    insert!(
        vars,
        "CONDA_DEFAULT_ENV",
        directories.host_prefix.to_string_lossy()
    );
    insert!(vars, "PREFIX", directories.host_prefix.to_string_lossy());
    insert!(
        vars,
        "BUILD_PREFIX",
        directories.build_prefix.to_string_lossy()
    );
    insert!(vars, "RECIPE_DIR", directories.recipe_dir.to_string_lossy());
    insert!(vars, "SRC_DIR", directories.work_dir.to_string_lossy());
    insert!(vars, "BUILD_DIR", directories.build_dir.to_string_lossy());
    insert!(
        vars,
        consts::RATTLER_BUILD_PACKAGE_FILES,
        directories.package_files_list_path().to_string_lossy()
    );

    // python variables
    // hard-code this because we never want pip's build isolation
    // https://github.com/conda/conda-build/pull/2972#discussion_r198290241
    //
    // Note that pip env "NO" variables are inverted logic.
    //    PIP_NO_BUILD_ISOLATION=False means don't use build isolation.
    insert!(vars, "PIP_NO_BUILD_ISOLATION", "False");
    // Some other env vars to have pip ignore dependencies. We supply them ourselves instead.
    insert!(vars, "PIP_NO_DEPENDENCIES", "True");
    insert!(vars, "PIP_IGNORE_INSTALLED", "True");

    // pip's cache directory (PIP_NO_CACHE_DIR) should not be
    // disabled as this results in .egg-info rather than
    // .dist-info directories being created, see gh-3094
    // set PIP_CACHE_DIR to a path in the work dir that does not exist.
    if let Some(work_dir_parent) = directories.work_dir.parent() {
        let pip_cache = work_dir_parent.join("pip_cache");
        insert!(vars, "PIP_CACHE_DIR", pip_cache.to_string_lossy());
    }

    // tell pip to not get anything from PyPI, please. We have everything we need
    // locally, and if we don't, it's a problem.
    insert!(vars, "PIP_NO_INDEX", "True");

    // For noarch packages, do not write any bytecode
    if output.is_python_version_independent() {
        insert!(vars, "PYTHONDONTWRITEBYTECODE", "1");
    }

    if output.build_configuration.force_colors {
        vars.extend(force_color_vars());
    }

    // pkg vars
    insert!(vars, "PKG_NAME", output.name().as_normalized());
    insert!(vars, "PKG_VERSION", output.version());
    insert!(
        vars,
        "PKG_BUILDNUM",
        output.recipe.build().number.unwrap_or(0).to_string()
    );

    let hash = output.build_configuration.hash.clone();
    insert!(vars, "PKG_BUILD_STRING", output.build_string().to_string());
    insert!(vars, "PKG_HASH", hash);

    if output.build_configuration.cross_compilation() {
        insert!(vars, "CONDA_BUILD_CROSS_COMPILATION", "1");
    } else {
        insert!(vars, "CONDA_BUILD_CROSS_COMPILATION", "0");
    }
    insert!(vars, "SUBDIR", output.target_platform().to_string());
    insert!(
        vars,
        "build_platform",
        output
            .build_configuration
            .build_platform
            .platform
            .to_string()
    );
    insert!(
        vars,
        "target_platform",
        output.target_platform().to_string()
    );
    insert!(
        vars,
        "host_platform",
        output.host_platform().platform.to_string()
    );
    insert!(vars, "CONDA_BUILD_STATE", build_state);

    vars.extend(language_vars(output));

    // for reproducibility purposes, set the SOURCE_DATE_EPOCH to the configured timestamp
    // this value will be taken from the previous package for rebuild purposes
    let timestamp_epoch_secs = output.build_configuration.timestamp.as_second();
    insert!(vars, "SOURCE_DATE_EPOCH", timestamp_epoch_secs);

    vars
}

/// Return environment variables that should be set during the test process.
///
/// This is a minimal subset of [`vars`] appropriate for test scripts — it does
/// not include build-time paths (`RECIPE_DIR`, `SRC_DIR`, …) or pip settings.
pub fn test_vars(
    target_platform: Platform,
    build_platform: Platform,
    host_platform: Platform,
) -> HashMap<String, Option<String>> {
    let mut vars = HashMap::new();

    insert!(vars, "CONDA_BUILD_STATE", "TEST");
    insert!(vars, "SUBDIR", target_platform.to_string());
    insert!(vars, "target_platform", target_platform.to_string());
    insert!(vars, "build_platform", build_platform.to_string());
    insert!(vars, "host_platform", host_platform.to_string());

    if let Some((_, arch)) = target_platform.to_string().rsplit_once('-') {
        insert!(vars, "ARCH", arch);
    }

    if target_platform != build_platform {
        insert!(vars, "CONDA_BUILD_CROSS_COMPILATION", "1");
    } else {
        insert!(vars, "CONDA_BUILD_CROSS_COMPILATION", "0");
    }

    vars
}

/// Language-specific variant keys that are excluded from environment variables
/// because they are handled separately (e.g. via `python_vars_from_records`).
const LANGUAGE_VARIANT_KEYS: &[&str] = &["PERL", "LUA", "R", "NUMPY", "PYTHON", "RUBY", "NODEJS"];

/// Convert a variant map into environment variables, excluding language-specific
/// keys that are handled separately.
pub fn env_vars_from_variant(
    variant: &BTreeMap<NormalizedKey, Variable>,
) -> HashMap<String, Option<String>> {
    let languages: HashSet<&str> = HashSet::from_iter(LANGUAGE_VARIANT_KEYS.iter().copied());
    variant
        .iter()
        .filter_map(|(k, v)| {
            let key_upper = k.normalize().to_uppercase();
            if !languages.contains(key_upper.as_str()) {
                Some((k.normalize(), Some(v.to_string())))
            } else {
                None
            }
        })
        .collect()
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn python_vars_use_site_packages_path_from_repodata() {
        use rattler_conda_types::{PackageName, PackageRecord, VersionWithSource};
        use std::str::FromStr;
        use url::Url;

        let mut package_record = PackageRecord::new(
            PackageName::from_str("python").unwrap(),
            "3.13.1".parse::<VersionWithSource>().unwrap(),
            "cp313t".to_string(),
        );
        package_record.python_site_packages_path =
            Some("lib/python3.13t/site-packages".to_string());
        let record = RepoDataRecord {
            package_record,
            identifier: "python-3.13.1-cp313t.conda".parse().unwrap(),
            url: Url::parse("https://example.com/python-3.13.1-cp313t.conda").unwrap(),
            channel: None,
        };

        let prefix = Path::new("prefix");
        let vars = python_vars_from_records(&[record], prefix, Platform::Linux64);

        assert_eq!(
            vars.get("SP_DIR")
                .and_then(|value| value.as_deref())
                .map(Path::new),
            Some(prefix.join("lib/python3.13t/site-packages").as_path())
        );
    }

    #[test]
    fn os_vars_uses_the_injected_runtime_environment() {
        let runtime = RuntimeEnv::for_test(Platform::Linux64).with_var("PATH", "/injected/bin");

        let vars = os_vars(
            Path::new("/some/prefix"),
            &Platform::Linux64,
            &Platform::Linux64,
            &Platform::Linux64,
            EnvironmentIsolation::CondaBuild,
            Path::new("/some/work"),
            &runtime,
        );

        assert_eq!(
            vars.get("PATH").and_then(|value| value.as_deref()),
            Some("/injected/bin")
        );
        assert_eq!(vars.get("LANG"), Some(&None));
    }

    /// noarch on a Windows host still emits Windows target vars (issue #2475).
    #[test]
    fn test_noarch_uses_host_platform_for_os_vars() {
        let prefix = Path::new("/some/prefix");
        let work_dir = Path::new("/some/work");

        let vars = os_vars(
            prefix,
            &Platform::NoArch,
            &Platform::Win64,
            &Platform::Win64,
            EnvironmentIsolation::Strict,
            work_dir,
            &RuntimeEnv::for_test(Platform::Win64),
        );

        assert!(vars.contains_key("LIBRARY_PREFIX"));
        assert!(vars.contains_key("LIBRARY_BIN"));
        assert!(vars.contains_key("SCRIPTS"));
        assert!(vars.contains_key("Path"));
        assert_eq!(
            vars.get("SHLIB_EXT").and_then(|v| v.as_deref()),
            Some(".dll")
        );
    }

    #[test]
    fn build_vars_follow_configured_build_platform() {
        let vars = os_vars(
            Path::new("/some/prefix"),
            &Platform::WinArm64,
            &Platform::WinArm64,
            &Platform::WinArm64,
            EnvironmentIsolation::Strict,
            Path::new("/some/work"),
            &RuntimeEnv::for_test(Platform::WinArm64),
        );

        assert_eq!(
            vars.get("BUILD").and_then(|value| value.as_deref()),
            Some("arm64-pc-windows-19.0.0")
        );
    }

    /// noarch on a non-Windows host does not emit Windows target vars.
    #[test]
    fn test_noarch_non_windows_host_has_no_windows_vars() {
        let prefix = Path::new("/some/prefix");
        let work_dir = Path::new("/some/work");

        let vars = os_vars(
            prefix,
            &Platform::NoArch,
            &Platform::Linux64,
            &Platform::Linux64,
            EnvironmentIsolation::Strict,
            work_dir,
            &RuntimeEnv::for_test(Platform::Linux64),
        );

        assert!(!vars.contains_key("LIBRARY_PREFIX"));
        assert!(vars.contains_key("PATH"));
    }
}
