use std::path::{Path, PathBuf};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("cargo:rerun-if-changed=proto/narsil/inference/v1/inference.proto");
    link_libtorch_cuda_if_present();

    let protoc = protoc_bin_vendored::protoc_bin_path()?;
    // Build scripts are single-threaded here; set PROTOC before prost-build
    // invokes protoc discovery.
    unsafe {
        std::env::set_var("PROTOC", protoc);
    }

    tonic_prost_build::configure()
        .compile_protos(&["proto/narsil/inference/v1/inference.proto"], &["proto"])?;

    Ok(())
}

fn link_libtorch_cuda_if_present() {
    println!("cargo:rustc-check-cfg=cfg(narsil_libtorch_cuda_linked)");
    println!("cargo:rerun-if-env-changed=LIBTORCH");
    println!("cargo:rerun-if-env-changed=LIBTORCH_LIB");
    println!("cargo:rerun-if-env-changed=LIBTORCH_USE_PYTORCH");
    println!("cargo:rerun-if-env-changed=PYTHON");

    if std::env::var_os("CARGO_FEATURE_TORCH").is_none() {
        return;
    }
    if std::env::var("CARGO_CFG_TARGET_OS").as_deref() != Ok("linux") {
        return;
    }

    let Some(lib_dir) = find_libtorch_lib_dir() else {
        return;
    };
    let cuda_lib = lib_dir.join("libtorch_cuda.so");
    if !cuda_lib.exists() {
        return;
    }
    // Re-run if a CPU libtorch is swapped for a CUDA one at the same lib dir
    // (the presence of libtorch_cuda.so is what flips narsil_libtorch_cuda_linked).
    println!("cargo:rerun-if-changed={}", cuda_lib.display());

    println!("cargo:rustc-link-search=native={}", lib_dir.display());
    println!("cargo:rustc-link-arg=-Wl,--push-state,--no-as-needed");
    println!("cargo:rustc-link-arg=-ltorch_cuda");
    println!("cargo:rustc-link-arg=-Wl,--pop-state");
    println!("cargo:rustc-cfg=narsil_libtorch_cuda_linked");
}

fn find_libtorch_lib_dir() -> Option<PathBuf> {
    // NOTE: `torch-sys` (which `tch` links via `links = "torch"`) exports its
    // resolved lib dir as `DEP_TORCH_LIBTORCH_LIB`, but Cargo only injects
    // `DEP_<links>_*` vars into the build script of a package's *direct*
    // dependents. narsil-engine depends on `tch`, not `torch-sys` directly, so
    // that var never reaches this build script — hence we re-derive the dir from
    // the same env knobs (`LIBTORCH_LIB` / `LIBTORCH` / `LIBTORCH_USE_PYTORCH`)
    // that torch-sys itself consults.
    if let Some(path) = std::env::var_os("LIBTORCH_LIB").map(PathBuf::from) {
        return Some(path);
    }
    if let Some(path) = std::env::var_os("LIBTORCH").map(PathBuf::from) {
        return Some(libtorch_lib_from_root(&path));
    }
    if std::env::var_os("LIBTORCH_USE_PYTORCH").is_some() {
        return python_libtorch_lib_dir();
    }
    None
}

fn libtorch_lib_from_root(root: &Path) -> PathBuf {
    let lib = root.join("lib");
    if lib.exists() {
        lib
    } else {
        root.to_path_buf()
    }
}

fn python_libtorch_lib_dir() -> Option<PathBuf> {
    // This branch is Linux-only and only matters for PyTorch 2.x (Python 3), so
    // default to `python3`: many distros (Ubuntu/Debian/Arch) ship no bare
    // `python`, and falling back to it would break auto-detection.
    let python = std::env::var_os("PYTHON").unwrap_or_else(|| "python3".into());
    let script = r#"
import pathlib
import torch
print(pathlib.Path(torch.__file__).resolve().parent / "lib")
"#;
    let output = std::process::Command::new(python)
        .arg("-c")
        .arg(script)
        .output()
        .ok()?;
    if !output.status.success() {
        return None;
    }
    let path = String::from_utf8(output.stdout).ok()?;
    let path = path.trim();
    (!path.is_empty()).then(|| PathBuf::from(path))
}
