athena_rs 3.6.1

use anyhow::{Context, Result};
use dirs::cache_dir;
use flate2::read::GzDecoder;
use once_cell::sync::OnceCell;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::time::Duration;
use tar::Archive;
use tokio::process::Command;
use tokio::task::spawn_blocking;
use tokio::{fs as tokio_fs, io::AsyncWriteExt};
use which::which;

const PG_VERSION: &str = "17.7";
const DEFAULT_PG_LINUX_X64_URL: &str =
    "https://get.enterprisedb.com/postgresql/postgresql-17.7-1-linux-x64-binaries.tar.gz";
const PG_ARCHIVE_NAME: &str = "postgresql-17.7-1-linux-x64-binaries.tar.gz";

static TOOLS_CELL: OnceCell<PgToolsPaths> = OnceCell::new();

/// Tracks whether we've already warned about foreign-OS `ATHENA_PG_*_PATH`
/// values so the operator only sees one warning per process instead of one
/// per resolution attempt (env_overrides is called twice: once before loading
/// `.env` and once after).
static FOREIGN_OS_WARNED: OnceCell<()> = OnceCell::new();

/// Paths to pg_dump and pg_restore binaries.
#[derive(Clone, Debug)]
pub struct PgToolsPaths {
    pub pg_dump: PathBuf,
    pub pg_restore: PathBuf,
}

fn tool_filename(base: &str) -> String {
    #[cfg(target_os = "windows")]
    {
        match base {
            "pg_dump" => "pg_dump.exe".to_string(),
            "pg_restore" => "pg_restore.exe".to_string(),
            other => other.to_string(),
        }
    }
    #[cfg(not(target_os = "windows"))]
    {
        base.to_string()
    }
}

fn parse_major_version(version: &str) -> Option<u32> {
    version.split('.').next()?.parse().ok()
}

fn parse_version_token(token: &str) -> Option<u32> {
    // Skip leading non-numeric characters until the first digit (e.g. "v17.7"),
    // then extract the numeric version portion before any suffix like "-ubuntu".
    let trimmed: &str = token.trim_start_matches(|ch: char| !ch.is_ascii_digit());
    let numeric: String = trimmed
        .chars()
        .take_while(|ch| ch.is_ascii_digit() || *ch == '.')
        .collect();
    if numeric.is_empty() {
        None
    } else {
        parse_major_version(&numeric)
    }
}

fn parse_pg_dump_version_major(output: &str) -> Option<u32> {
    output.split_whitespace().find_map(parse_version_token)
}

async fn pg_dump_major_version_internal(path: &Path) -> Option<u32> {
    let output: std::process::Output = match Command::new(path).arg("--version").output().await {
        Ok(output) => output,
        Err(err) => {
            tracing::warn!(
                "Failed to execute pg_dump at {} (path may not exist or is not executable): {}",
                path.display(),
                err
            );
            return None;
        }
    };
    let stdout: std::borrow::Cow<'_, str> = String::from_utf8_lossy(&output.stdout);
    let stderr: std::borrow::Cow<'_, str> = String::from_utf8_lossy(&output.stderr);
    let text: std::borrow::Cow<'_, str> = if stdout.trim().is_empty() {
        stderr
    } else {
        stdout
    };
    let parsed = parse_pg_dump_version_major(&text);
    if parsed.is_none() {
        tracing::warn!(
            "Unable to parse pg_dump version output from {}: {}",
            path.display(),
            text.trim()
        );
    }
    parsed
}

/// Return the required pg_dump major version for this build.
///
/// Athena’s backup/restore endpoints shell out to `pg_dump` / `pg_restore`. A
/// pg_dump major version that is **older** than the server will fail with a
/// version mismatch; using a newer major is safe.
pub fn required_pg_dump_major() -> u32 {
    parse_major_version(PG_VERSION).unwrap_or(17)
}

/// Read `pg_dump --version` and return the parsed major version.
pub async fn pg_dump_major_version(path: &Path) -> Option<u32> {
    pg_dump_major_version_internal(path).await
}

/// Resolve `pg_dump` and `pg_restore` from a versioned tools directory.
///
/// Directory layout:
/// - `$ATHENA_PG_TOOLS_DIR/<major>/bin/pg_dump`
/// - `$ATHENA_PG_TOOLS_DIR/<major>/bin/pg_restore`
///
/// This enables supporting multiple Postgres server majors on the same host by
/// installing multiple client tool versions side-by-side and selecting the
/// appropriate one at runtime.
pub fn resolve_pg_tools_from_dir(server_major: u32) -> Option<PgToolsPaths> {
    let root: String = std::env::var("ATHENA_PG_TOOLS_DIR").ok()?;
    if root.trim().is_empty() {
        return None;
    }

    let bin_dir: PathBuf = PathBuf::from(root)
        .join(server_major.to_string())
        .join("bin");

    let dump: PathBuf = bin_dir.join(tool_filename("pg_dump"));
    let restore: PathBuf = bin_dir.join(tool_filename("pg_restore"));

    if dump.is_file() && restore.is_file() {
        Some(PgToolsPaths {
            pg_dump: dump,
            pg_restore: restore,
        })
    } else {
        None
    }
}

/// Resolve `pg_dump` / `pg_restore` from `$ATHENA_PG_TOOLS_DIR` without knowing
/// a specific server major.
///
/// Preference order:
/// 1. The required major for this build (`required_pg_dump_major()`).
/// 2. The highest integer-named `<major>/bin/` directory present that contains
///    both binaries — newer client majors can dump older servers safely.
fn resolve_pg_tools_from_dir_best_effort() -> Option<PgToolsPaths> {
    let root_raw: String = std::env::var("ATHENA_PG_TOOLS_DIR").ok()?;
    if root_raw.trim().is_empty() {
        return None;
    }
    let root: PathBuf = PathBuf::from(root_raw.trim());

    if let Some(paths) = resolve_pg_tools_from_dir(required_pg_dump_major()) {
        return Some(paths);
    }

    let read_dir: std::fs::ReadDir = std::fs::read_dir(&root).ok()?;
    let mut majors: Vec<u32> = read_dir
        .filter_map(|entry| entry.ok())
        .filter(|entry| entry.file_type().map(|ft| ft.is_dir()).unwrap_or(false))
        .filter_map(|entry| {
            entry
                .file_name()
                .to_str()
                .and_then(|name| name.parse::<u32>().ok())
        })
        .collect();
    majors.sort_unstable_by(|a, b| b.cmp(a));

    majors.into_iter().find_map(resolve_pg_tools_from_dir)
}

/// Return `Some(downloaded_paths)` when the installed pg_dump is older than
/// required and a newer bundle can be downloaded.
/// Returns `None` when the installed version is sufficient or download fails.
#[allow(dead_code)]
async fn maybe_download_newer_pg_tools(dump_path: &Path) -> Option<PgToolsPaths> {
    let required_major: u32 = parse_major_version(PG_VERSION)?;
    let installed_major: u32 = pg_dump_major_version_internal(dump_path).await?;

    if installed_major >= required_major {
        return None;
    }

    tracing::info!(
        "pg_dump on PATH is major version {} which is older than required major {} (PG_VERSION {}). Attempting to download pg tools.",
        installed_major,
        required_major,
        PG_VERSION
    );

    if !allow_pg_tools_download() {
        tracing::warn!(
            "Automatic pg tools download disabled; using existing pg_dump major version {}.",
            installed_major
        );
        return None;
    }

    if let Err(err) = ensure_linux_x64() {
        tracing::warn!(
            "Automatic pg tools download not available: {}. Using existing pg_dump major version {}.",
            err,
            installed_major
        );
        return None;
    }

    match download_and_extract().await {
        Ok(paths) => Some(paths),
        Err(err) => {
            tracing::warn!(
                "Failed to download newer pg tools: {}. Using existing pg_dump major version {}.",
                err,
                installed_major
            );
            None
        }
    }
}

/// Ensure pg_dump and pg_restore are available, downloading a portable bundle if needed.
///
/// Resolution order:
/// 1) `ATHENA_PG_DUMP_PATH` / `ATHENA_PG_RESTORE_PATH`
/// 2) `$ATHENA_PG_TOOLS_DIR/<major>/bin/{pg_dump,pg_restore}` — prefers the
///    required major, falls back to the highest installed major
/// 3) Binaries found on PATH or predefined install path on Windows
/// 4) Download Linux x86_64 portable client bundle into cache (~/.cache/athena/pg_tools/<ver>)
pub async fn ensure_pg_tools() -> Result<PgToolsPaths> {
    if let Some(cached) = TOOLS_CELL.get() {
        return Ok(cached.clone());
    }

    // 1) Explicit env overrides from current environment.
    if let Some(paths) = env_overrides()? {
        let _ = TOOLS_CELL.set(paths.clone());
        return Ok(paths);
    }

    // 1.5) Try to read from .env if not already set
    {
        // Load .env if present and re-try env_overrides
        let dotenv_path: Option<PathBuf> = std::env::current_dir()
            .ok()
            .map(|dir| dir.join(".env"))
            .filter(|path| path.exists());
        if let Some(dotenv_file) = dotenv_path {
            // Load .env into process env once; ignore errors if file is malformed or missing.
            let _ = dotenv::from_path(&dotenv_file);
            // After loading .env, re-check for explicit env overrides
            if let Some(paths) = env_overrides()? {
                let _ = TOOLS_CELL.set(paths.clone());
                return Ok(paths);
            }
        }
    }

    // 2) ATHENA_PG_TOOLS_DIR layout (populated by scripts/provisioning/10-setup-pg-tools-dir.sh).
    if let Some(paths) = resolve_pg_tools_from_dir_best_effort() {
        tracing::info!(
            pg_dump = %paths.pg_dump.display(),
            pg_restore = %paths.pg_restore.display(),
            "Resolved pg tools from ATHENA_PG_TOOLS_DIR"
        );
        let _ = TOOLS_CELL.set(paths.clone());
        return Ok(paths);
    }

    // 2) PATH lookup or fallback on Windows to known install location.
    #[cfg(target_os = "windows")]
    {
        let candidates = [("pg_dump", "pg_dump.exe"), ("pg_restore", "pg_restore.exe")];
        let which_with_exe = |tool_name: &str, exe_name: &str| -> Option<PathBuf> {
            which(exe_name).ok().or_else(|| which(tool_name).ok())
        };

        let dump_opt: Option<PathBuf> = which_with_exe(candidates[0].0, candidates[0].1);
        let restore_opt: Option<PathBuf> = which_with_exe(candidates[1].0, candidates[1].1);

        // If found on PATH, use those
        if let (Some(dump), Some(restore)) = (dump_opt.clone(), restore_opt.clone()) {
            let paths: PgToolsPaths = PgToolsPaths {
                pg_dump: dump,
                pg_restore: restore,
            };
            let _ = TOOLS_CELL.set(paths.clone());
            return Ok(paths);
        }

        // Otherwise, try C:\Program Files\PostgreSQL\18\bin\
        let dump_path: PathBuf = PathBuf::from(r"C:\Program Files\PostgreSQL\18\bin\pg_dump.exe");
        let restore_path: PathBuf =
            PathBuf::from(r"C:\Program Files\PostgreSQL\18\bin\pg_restore.exe");
        if dump_path.is_file() && restore_path.is_file() {
            let paths = PgToolsPaths {
                pg_dump: dump_path,
                pg_restore: restore_path,
            };
            let _ = TOOLS_CELL.set(paths.clone());
            return Ok(paths);
        }
    }

    #[cfg(not(target_os = "windows"))]
    {
        let candidates = [("pg_dump", "pg_dump"), ("pg_restore", "pg_restore")];
        let which_with_exe =
            |tool_name: &str, _exe_name: &str| -> Option<PathBuf> { which(tool_name).ok() };

        // First attempt: use whatever is already installed on PATH.
        let mut dump_opt: Option<PathBuf> = which_with_exe(candidates[0].0, candidates[0].1);
        let mut restore_opt: Option<PathBuf> = which_with_exe(candidates[1].0, candidates[1].1);

        if let (Some(dump), Some(restore)) = (dump_opt.clone(), restore_opt.clone()) {
            if let Some(paths) = maybe_download_newer_pg_tools(&dump).await {
                if TOOLS_CELL.set(paths.clone()).is_err() {
                    // Another caller already initialized the cache; prefer the cached paths
                    // so all requests share the same resolved tools.
                    if let Some(cached) = TOOLS_CELL.get() {
                        return Ok(cached.clone());
                    }
                    tracing::warn!(
                        "pg tool cache already initialized but no cached value found; using freshly downloaded tools."
                    );
                }
                return Ok(paths);
            }
            let paths: PgToolsPaths = PgToolsPaths {
                pg_dump: dump,
                pg_restore: restore,
            };
            let _ = TOOLS_CELL.set(paths.clone());
            return Ok(paths);
        }

        // Second attempt (Linux only): try to install PostgreSQL client tools via apt-get,
        // then re-check PATH. This is best-effort and will be skipped on non-Debian distros
        // or when apt-get is unavailable.
        #[cfg(target_os = "linux")]
        {
            maybe_install_pg_tools_with_apt().await;

            dump_opt = which_with_exe(candidates[0].0, candidates[0].1);
            restore_opt = which_with_exe(candidates[1].0, candidates[1].1);

            if let (Some(dump), Some(restore)) = (dump_opt, restore_opt) {
                let paths: PgToolsPaths = PgToolsPaths {
                    pg_dump: dump,
                    pg_restore: restore,
                };
                let _ = TOOLS_CELL.set(paths.clone());
                return Ok(paths);
            }
        }
    }

    // 4) Download portable bundle (Linux x86_64 only), unless explicitly disabled.
    if !allow_pg_tools_download() {
        anyhow::bail!(
            "pg_dump/pg_restore not found on PATH, not resolvable from ATHENA_PG_TOOLS_DIR, \
             and automatic download is disabled. Install PostgreSQL client tools \
             (e.g. apt install postgresql-client postgresql-common), populate \
             $ATHENA_PG_TOOLS_DIR/<major>/bin/ (see scripts/provisioning/10-setup-pg-tools-dir.sh), \
             set ATHENA_PG_DUMP_PATH and ATHENA_PG_RESTORE_PATH, or use the official \
             Athena Docker image which has them pre-installed."
        );
    }
    ensure_linux_x64()?;
    let paths: PgToolsPaths = download_and_extract().await?;
    let _ = TOOLS_CELL.set(paths.clone());
    Ok(paths)
}

/// When true, we may attempt to download the portable pg_dump bundle from the network.
/// Set ATHENA_PG_TOOLS_ALLOW_DOWNLOAD=0 or false to disable (e.g. in Docker with tools pre-installed).
fn allow_pg_tools_download() -> bool {
    match std::env::var("ATHENA_PG_TOOLS_ALLOW_DOWNLOAD") {
        Ok(v) => !matches!(v.as_str(), "0" | "false" | "FALSE" | "False" | "no"),
        Err(_) => true,
    }
}

/// Returns `true` when `raw` looks like an absolute path that belongs to the
/// current operating system.
///
/// Used to gracefully ignore env-var overrides that leaked across platforms
/// — for example a Windows path like `C:/Program Files/PostgreSQL/18/bin/pg_dump.exe`
/// that ended up in a `.env` file consumed by a Linux container via
/// `docker-compose`'s `env_file:`.
///
/// Rules:
/// - Windows: drive-letter prefix (`C:...`) or UNC prefix (`\\server\share`).
///   POSIX-absolute paths like `/usr/bin/pg_dump` are rejected.
/// - Non-Windows: any path that does not look Windows-shaped is accepted,
///   including relative paths (e.g. `pg_dump`).
fn path_belongs_to_current_os(raw: &str) -> bool {
    let s: &str = raw.trim();
    if s.is_empty() {
        return false;
    }
    let bytes: &[u8] = s.as_bytes();
    let has_drive_letter: bool =
        bytes.len() >= 2 && bytes[0].is_ascii_alphabetic() && bytes[1] == b':';
    let has_unc_prefix: bool = s.starts_with(r"\\");
    let looks_windows: bool = has_drive_letter || has_unc_prefix;
    let looks_posix_absolute: bool = s.starts_with('/') || s.starts_with("~/");

    if cfg!(target_os = "windows") {
        // Accept drive-letter, UNC, and relative paths; reject obviously
        // POSIX-only absolute paths such as "/usr/bin/pg_dump".
        !(looks_posix_absolute && !looks_windows)
    } else {
        !looks_windows
    }
}

/// Candidate install roots to scan for `pg_dump` / `pg_restore`, per OS.
///
/// Globs like `/usr/lib/postgresql/*/bin` are expanded by walking the parent
/// directory and matching numeric major-version subdirectories (newer
/// versions first).
fn candidate_install_roots() -> Vec<PathBuf> {
    let mut roots: Vec<PathBuf> = Vec::new();

    #[cfg(target_os = "linux")]
    {
        // Debian/Ubuntu PGDG layout (postgresql-client-<major>).
        roots.extend(expand_versioned_dirs(
            Path::new("/usr/lib/postgresql"),
            "bin",
        ));
        // RHEL/Fedora/Rocky PGDG layout (postgresql-server / pgdg-<major>).
        roots.extend(expand_versioned_dirs(
            Path::new("/usr/pgsql-"),
            /* child */ "bin",
        ));
        // Generic fallbacks.
        roots.push(PathBuf::from("/usr/bin"));
        roots.push(PathBuf::from("/usr/local/bin"));
    }

    #[cfg(target_os = "macos")]
    {
        // Homebrew (Apple Silicon and Intel) versioned kegs.
        for prefix in ["/opt/homebrew/opt", "/usr/local/opt"] {
            let prefix_path: &Path = Path::new(prefix);
            if let Ok(read_dir) = std::fs::read_dir(prefix_path) {
                let mut pg_dirs: Vec<(u32, PathBuf)> = read_dir
                    .filter_map(|entry| entry.ok())
                    .filter_map(|entry| {
                        let name: String = entry.file_name().to_string_lossy().into_owned();
                        let major_str: &str = name.strip_prefix("postgresql@")?;
                        let major: u32 = major_str.parse().ok()?;
                        Some((major, entry.path().join("bin")))
                    })
                    .collect();
                pg_dirs.sort_unstable_by(|a, b| b.0.cmp(&a.0));
                roots.extend(pg_dirs.into_iter().map(|(_, p)| p));
            }
        }
        // EnterpriseDB installer layout: /Library/PostgreSQL/<major>/bin.
        roots.extend(expand_versioned_dirs(
            Path::new("/Library/PostgreSQL"),
            "bin",
        ));
        // Generic fallbacks.
        roots.push(PathBuf::from("/opt/homebrew/bin"));
        roots.push(PathBuf::from("/usr/local/bin"));
        roots.push(PathBuf::from("/usr/bin"));
    }

    #[cfg(target_os = "windows")]
    {
        // EnterpriseDB installer default layout (Program Files and 32-bit mirror).
        for prefix in [
            r"C:\Program Files\PostgreSQL",
            r"C:\Program Files (x86)\PostgreSQL",
        ] {
            roots.extend(expand_versioned_dirs(Path::new(prefix), "bin"));
        }
    }

    roots
}

/// Expand `root/<numeric_major>/<child>` directories, newest major first.
///
/// Returns an empty vec when `root` does not exist. Non-numeric subdirectory
/// names (e.g. `common`, `client`) are ignored.
fn expand_versioned_dirs(root: &Path, child: &str) -> Vec<PathBuf> {
    let Ok(read_dir) = std::fs::read_dir(root) else {
        return Vec::new();
    };
    let mut pairs: Vec<(u32, PathBuf)> = read_dir
        .filter_map(|entry| entry.ok())
        .filter_map(|entry| {
            let ft: std::fs::FileType = entry.file_type().ok()?;
            if !ft.is_dir() {
                return None;
            }
            let name: String = entry.file_name().to_string_lossy().into_owned();
            let major: u32 = name.parse().ok()?;
            Some((major, entry.path().join(child)))
        })
        .collect();
    pairs.sort_unstable_by(|a, b| b.0.cmp(&a.0));
    pairs.into_iter().map(|(_, p)| p).collect()
}

/// Actively search the current OS for `pg_dump` / `pg_restore`.
///
/// Strategy (first hit wins):
/// 1. `which pg_dump` / `which pg_restore` (PATH lookup).
/// 2. OS-specific install roots returned by [`candidate_install_roots`]
///    (e.g. `/usr/lib/postgresql/<major>/bin` on Debian, Homebrew kegs on
///    macOS, `C:\Program Files\PostgreSQL\<major>\bin` on Windows).
///
/// Both binaries must resolve to existing files for `Some(_)` to be returned.
fn discover_pg_tools_for_current_os() -> Option<PgToolsPaths> {
    let dump_name: String = tool_filename("pg_dump");
    let restore_name: String = tool_filename("pg_restore");

    if let (Ok(dump), Ok(restore)) = (which(&dump_name), which(&restore_name)) {
        if dump.is_file() && restore.is_file() {
            return Some(PgToolsPaths {
                pg_dump: dump,
                pg_restore: restore,
            });
        }
    }

    for root in candidate_install_roots() {
        let dump: PathBuf = root.join(&dump_name);
        let restore: PathBuf = root.join(&restore_name);
        if dump.is_file() && restore.is_file() {
            return Some(PgToolsPaths {
                pg_dump: dump,
                pg_restore: restore,
            });
        }
    }

    None
}

/// Replace the `ATHENA_PG_DUMP_PATH` / `ATHENA_PG_RESTORE_PATH` env vars with
/// `paths` so downstream code (and subprocesses we spawn) see the
/// newly-discovered, OS-appropriate locations.
///
/// # Safety
/// `std::env::set_var` is `unsafe` since Rust 1.80 because it is not thread
/// safe on POSIX. This is called during early startup (`ensure_pg_tools`)
/// before worker threads are spawned, which is the documented safe window.
fn rewrite_pg_tools_env(paths: &PgToolsPaths) {
    let dump: &str = &paths.pg_dump.to_string_lossy();
    let restore: &str = &paths.pg_restore.to_string_lossy();
    unsafe {
        std::env::set_var("ATHENA_PG_DUMP_PATH", dump);
        std::env::set_var("ATHENA_PG_RESTORE_PATH", restore);
    }
}

fn env_overrides() -> Result<Option<PgToolsPaths>> {
    let dump_raw: Option<String> = std::env::var("ATHENA_PG_DUMP_PATH").ok();
    let restore_raw: Option<String> = std::env::var("ATHENA_PG_RESTORE_PATH").ok();

    let current_os: &str = std::env::consts::OS;

    let dump_is_foreign: bool = dump_raw
        .as_deref()
        .map(|value| !value.trim().is_empty() && !path_belongs_to_current_os(value))
        .unwrap_or(false);
    let restore_is_foreign: bool = restore_raw
        .as_deref()
        .map(|value| !value.trim().is_empty() && !path_belongs_to_current_os(value))
        .unwrap_or(false);

    // When the operator set a foreign-OS override (classic case: a Windows
    // path in `.env` loaded into a Linux container via docker-compose's
    // `env_file:`), actively scan this host for pg_dump/pg_restore and
    // rewrite the env vars in-place so the rest of the app — and any
    // subprocesses — see the corrected paths.
    if dump_is_foreign || restore_is_foreign {
        if FOREIGN_OS_WARNED.set(()).is_ok() {
            if let Some(value) = dump_raw.as_deref().filter(|_| dump_is_foreign) {
                tracing::warn!(
                    path = %value,
                    os = %current_os,
                    "ATHENA_PG_DUMP_PATH does not look like a {current_os} path; searching system for a replacement"
                );
            }
            if let Some(value) = restore_raw.as_deref().filter(|_| restore_is_foreign) {
                tracing::warn!(
                    path = %value,
                    os = %current_os,
                    "ATHENA_PG_RESTORE_PATH does not look like a {current_os} path; searching system for a replacement"
                );
            }
        }

        if let Some(discovered) = discover_pg_tools_for_current_os() {
            tracing::info!(
                pg_dump = %discovered.pg_dump.display(),
                pg_restore = %discovered.pg_restore.display(),
                os = %current_os,
                "Auto-discovered pg_dump/pg_restore on {current_os}; rewriting ATHENA_PG_DUMP_PATH/ATHENA_PG_RESTORE_PATH"
            );
            rewrite_pg_tools_env(&discovered);
            return Ok(Some(discovered));
        }

        tracing::warn!(
            "No pg_dump/pg_restore found on {current_os} after scanning PATH and known install locations; falling back to ATHENA_PG_TOOLS_DIR / download"
        );
        // Fall through: treat both as unset and let later resolvers run.
        return Ok(None);
    }

    let dump_effective: Option<String> = dump_raw.filter(|v| !v.trim().is_empty());
    let restore_effective: Option<String> = restore_raw.filter(|v| !v.trim().is_empty());

    if dump_effective.is_none() && restore_effective.is_none() {
        return Ok(None);
    }

    let dump_path: PathBuf = dump_effective
        .map(PathBuf::from)
        .context("ATHENA_PG_DUMP_PATH set but empty or foreign-OS")?;
    let restore_path: PathBuf = restore_effective.map(PathBuf::from).unwrap_or_else(|| {
        dump_path
            .parent()
            .map(|p| p.join(tool_filename("pg_restore")))
            .unwrap_or_else(|| dump_path.clone())
    });

    if !dump_path.is_file() {
        anyhow::bail!("pg_dump not found at {}", dump_path.display());
    }
    if !restore_path.is_file() {
        anyhow::bail!("pg_restore not found at {}", restore_path.display());
    }

    Ok(Some(PgToolsPaths {
        pg_dump: dump_path,
        pg_restore: restore_path,
    }))
}

fn ensure_linux_x64() -> Result<()> {
    if cfg!(all(target_os = "linux", target_arch = "x86_64")) {
        Ok(())
    } else {
        anyhow::bail!(
            "Automatic pg_dump download is only supported on Linux x86_64; set ATHENA_PG_DUMP_PATH and ATHENA_PG_RESTORE_PATH instead"
        )
    }
}

async fn download_and_extract() -> Result<PgToolsPaths> {
    let cache_root: PathBuf = cache_dir()
        .unwrap_or(std::env::temp_dir())
        .join("athena")
        .join("pg_tools")
        .join(format!("{}-linux-x64", PG_VERSION));
    let bin_dir: PathBuf = cache_root.join("bin");
    let pg_dump_path: PathBuf = bin_dir.join("pg_dump");
    let pg_restore_path: PathBuf = bin_dir.join("pg_restore");

    if pg_dump_path.is_file() && pg_restore_path.is_file() {
        return Ok(PgToolsPaths {
            pg_dump: pg_dump_path,
            pg_restore: pg_restore_path,
        });
    }

    tokio_fs::create_dir_all(&cache_root)
        .await
        .context("create cache dir")?;
    tokio_fs::create_dir_all(&bin_dir).await?;

    let archive_path: PathBuf = cache_root.join(PG_ARCHIVE_NAME);
    if !archive_path.is_file() {
        let url = std::env::var("ATHENA_PG_TOOLS_URL")
            .unwrap_or_else(|_| DEFAULT_PG_LINUX_X64_URL.to_string());
        download_archive(&url, &archive_path).await?;
    }

    extract_binaries(&archive_path, &bin_dir).await?;

    Ok(PgToolsPaths {
        pg_dump: pg_dump_path,
        pg_restore: pg_restore_path,
    })
}

async fn download_archive(url: &str, dest: &Path) -> Result<()> {
    let client: reqwest::Client = reqwest::Client::builder()
        .timeout(Duration::from_secs(60))
        .build()?;

    let resp: reqwest::Response = client
        .get(url)
        .send()
        .await
        .with_context(|| format!("downloading {}", url))?;
    let status: reqwest::StatusCode = resp.status();
    if !status.is_success() {
        let hint = if status.as_u16() == 403 {
            " (403 Forbidden — use an image with PostgreSQL client pre-installed, or set ATHENA_PG_DUMP_PATH/ATHENA_PG_RESTORE_PATH)"
        } else {
            ""
        };
        anyhow::bail!("download failed ({}){}: {}", status, hint, url);
    }

    let bytes: axum::body::Bytes = resp.bytes().await?;
    let mut file: tokio_fs::File = tokio_fs::File::create(dest).await?;
    file.write_all(&bytes).await?;
    file.flush().await?;
    Ok(())
}

/// Best-effort installation of PostgreSQL client tools via `apt-get` on Linux.
///
/// This helper:
/// - no-ops when `apt-get` is not present (non-Debian distros, containers without apt)
/// - respects `ATHENA_AUTO_INSTALL_PG_TOOLS=false` / `0` to opt out
/// - logs failures but does not return an error (download fallback still applies)
#[cfg(target_os = "linux")]
async fn maybe_install_pg_tools_with_apt() {
    use tokio::process::Command;

    // Allow operators to explicitly disable this behavior.
    if let Ok(val) = std::env::var("ATHENA_AUTO_INSTALL_PG_TOOLS") {
        if matches!(val.as_str(), "0" | "false" | "FALSE" | "False") {
            tracing::info!(
                "ATHENA_AUTO_INSTALL_PG_TOOLS disabled; skipping automatic pg tools install via apt-get."
            );
            return;
        }
    }

    // Only attempt installation when apt-get is available.
    if which("apt-get").is_err() {
        tracing::info!("apt-get not found on PATH; skipping automatic pg tools install.");
        return;
    }

    tracing::info!(
        "Attempting to install PostgreSQL client tools via apt-get (postgresql-client, postgresql-common)."
    );

    // Run `apt-get update` first.
    match Command::new("apt-get").arg("update").status().await {
        Ok(status) if status.success() => {
            tracing::info!("apt-get update succeeded before installing PostgreSQL client tools.");
        }
        Ok(status) => {
            tracing::warn!(
                ?status,
                "apt-get update failed; skipping automatic pg tools install."
            );
            return;
        }
        Err(err) => {
            tracing::warn!(error = %err, "Failed to invoke apt-get update; skipping automatic pg tools install.");
            return;
        }
    }

    // Then try to install the generic PostgreSQL client packages.
    match Command::new("apt-get")
        .args(["install", "-y", "postgresql-client", "postgresql-common"])
        .status()
        .await
    {
        Ok(status) if status.success() => {
            tracing::info!(
                "apt-get install postgresql-client postgresql-common succeeded; pg_dump/pg_restore should now be on PATH."
            );
        }
        Ok(status) => {
            tracing::warn!(
                ?status,
                "apt-get install postgresql-client postgresql-common failed; pg tools may still be missing."
            );
        }
        Err(err) => {
            tracing::warn!(error = %err, "Failed to invoke apt-get install; pg tools may still be missing.");
        }
    }
}

async fn extract_binaries(archive_path: &Path, bin_dir: &Path) -> Result<()> {
    let archive_path: PathBuf = archive_path.to_owned();
    let bin_dir: PathBuf = bin_dir.to_owned();

    spawn_blocking(move || -> Result<()> {
        let file: File = File::open(&archive_path)?;
        let decoder: GzDecoder<File> = GzDecoder::new(file);
        let mut archive: Archive<GzDecoder<File>> = Archive::new(decoder);

        let mut extracted: usize = 0usize;
        for entry in archive.entries()? {
            let mut entry: tar::Entry<'_, GzDecoder<File>> = entry?;
            let path: std::borrow::Cow<'_, Path> = entry.path()?;
            let name: &str = path
                .file_name()
                .and_then(|v| v.to_str())
                .unwrap_or_default();
            if name == "pg_dump" || name == "pg_restore" {
                let dest = bin_dir.join(name);
                entry.unpack(&dest)?;
                #[cfg(unix)]
                {
                    let mut perms: std::fs::Permissions = std::fs::metadata(&dest)?.permissions();
                    use std::os::unix::fs::PermissionsExt;
                    perms.set_mode(0o755);
                    std::fs::set_permissions(&dest, perms)?;
                }
                extracted += 1;
            }
        }

        if extracted < 2 {
            anyhow::bail!("pg_dump/pg_restore not found in downloaded archive");
        }

        Ok(())
    })
    .await?
}

#[cfg(test)]
mod tests {
    use std::path::Path;

    use super::{
        maybe_download_newer_pg_tools, parse_major_version, parse_pg_dump_version_major,
        parse_version_token, path_belongs_to_current_os, pg_dump_major_version,
    };

    #[test]
    fn parse_pg_dump_version_major_from_standard_output() {
        let output: &str = "pg_dump (PostgreSQL) 17.7";
        assert_eq!(parse_pg_dump_version_major(output), Some(17));
    }

    #[test]
    fn parse_pg_dump_version_major_from_ubuntu_output() {
        let output = "pg_dump (PostgreSQL) 16.13 (Ubuntu 16.13-0ubuntu0.24.04.1)";
        assert_eq!(parse_pg_dump_version_major(output), Some(16));
    }

    #[test]
    fn parse_pg_dump_version_major_from_token_suffix() {
        let output: &str = "pg_dump (PostgreSQL) 16.13-0ubuntu0.24.04.1";
        assert_eq!(parse_pg_dump_version_major(output), Some(16));
    }

    #[test]
    fn parse_major_version_handles_invalid() {
        assert_eq!(parse_major_version(""), None);
        assert_eq!(parse_major_version("abc"), None);
        assert_eq!(parse_major_version("9"), Some(9));
        assert_eq!(parse_major_version("17.7"), Some(17));
    }

    #[test]
    fn parse_version_token_handles_prefix_and_suffix() {
        assert_eq!(parse_version_token("16.13-0ubuntu0.24.04.1"), Some(16));
        assert_eq!(parse_version_token("v17.7"), Some(17));
        assert_eq!(parse_version_token("PostgreSQL17.7"), Some(17));
        assert_eq!(parse_version_token("postgresql"), None);
    }

    #[tokio::test]
    async fn maybe_download_returns_none_when_version_check_fails() {
        let missing_path: &Path = Path::new("/nonexistent/pg_dump");
        assert!(maybe_download_newer_pg_tools(missing_path).await.is_none());
    }

    #[tokio::test]
    async fn pg_dump_major_version_parses_pg_dump_when_available() {
        let Ok(pg_dump_path) = which::which("pg_dump") else {
            return;
        };
        let parsed: Option<u32> = pg_dump_major_version(&pg_dump_path).await;
        assert!(parsed.is_some());
    }

    #[test]
    fn path_belongs_to_current_os_rejects_empty_and_whitespace() {
        assert!(!path_belongs_to_current_os(""));
        assert!(!path_belongs_to_current_os("   "));
    }

    #[cfg(not(target_os = "windows"))]
    #[test]
    fn path_belongs_to_current_os_on_unix() {
        // POSIX-absolute and relative accepted.
        assert!(path_belongs_to_current_os("/usr/bin/pg_dump"));
        assert!(path_belongs_to_current_os(
            "/usr/lib/postgresql/17/bin/pg_dump"
        ));
        assert!(path_belongs_to_current_os("~/pg/pg_dump"));
        assert!(path_belongs_to_current_os("pg_dump"));
        assert!(path_belongs_to_current_os("./bin/pg_dump"));

        // Windows-shaped paths rejected.
        assert!(!path_belongs_to_current_os(
            "C:/Program Files/PostgreSQL/18/bin/pg_dump.exe"
        ));
        assert!(!path_belongs_to_current_os(
            r"C:\Program Files\PostgreSQL\18\bin\pg_dump.exe"
        ));
        assert!(!path_belongs_to_current_os(r"\\server\share\pg_dump.exe"));
        assert!(!path_belongs_to_current_os("D:/pg_dump.exe"));
    }

    #[cfg(target_os = "windows")]
    #[test]
    fn path_belongs_to_current_os_on_windows() {
        // Drive-letter, UNC, and relative accepted.
        assert!(path_belongs_to_current_os(
            "C:/Program Files/PostgreSQL/18/bin/pg_dump.exe"
        ));
        assert!(path_belongs_to_current_os(
            r"C:\Program Files\PostgreSQL\18\bin\pg_dump.exe"
        ));
        assert!(path_belongs_to_current_os(r"\\server\share\pg_dump.exe"));
        assert!(path_belongs_to_current_os("pg_dump.exe"));
        assert!(path_belongs_to_current_os(r".\bin\pg_dump.exe"));

        // POSIX-only absolute paths rejected.
        assert!(!path_belongs_to_current_os("/usr/bin/pg_dump"));
        assert!(!path_belongs_to_current_os(
            "/usr/lib/postgresql/17/bin/pg_dump"
        ));
        assert!(!path_belongs_to_current_os("~/pg/pg_dump"));
    }

    #[test]
    fn resolve_pg_tools_from_dir_uses_expected_layout() {
        let tmp: std::path::PathBuf =
            std::env::temp_dir().join(format!("athena_pg_tools_test_{}", uuid::Uuid::new_v4()));
        let bin_dir: std::path::PathBuf = tmp.join("17").join("bin");
        std::fs::create_dir_all(&bin_dir).expect("create test bin dir");

        let dump: std::path::PathBuf = bin_dir.join(super::tool_filename("pg_dump"));
        let restore: std::path::PathBuf = bin_dir.join(super::tool_filename("pg_restore"));
        std::fs::write(&dump, b"").expect("create dump file");
        std::fs::write(&restore, b"").expect("create restore file");

        unsafe {
            std::env::set_var("ATHENA_PG_TOOLS_DIR", &tmp);
        }
        let resolved: crate::utils::pg_tools::PgToolsPaths =
            super::resolve_pg_tools_from_dir(17).expect("resolved tools");
        assert_eq!(resolved.pg_dump, dump);
        assert_eq!(resolved.pg_restore, restore);

        let _ = std::fs::remove_dir_all(&tmp);
    }
}