//! CUDA kernel precompilation utilities.
//!
//! Compiles `.cu` source to native cubin (SASS) via `nvcc` with full `-O3`
//! optimization and a deterministic SHA-256 disk cache. The cubin is
//! architecture-specific (`sm_XX`) and eliminates driver JIT entirely.
//!
//! Moved from `ml::cuda_pipeline` so that `ml-dqn` (which depends on `ml-core`
//! but not `ml`) can also precompile kernels without a circular dependency.

/// nvcc optimization flags applied to every kernel compilation.
///
/// Changing this array auto-invalidates the cubin cache (flags are hashed into
/// the cache key). No manual version bump needed.
///
/// Note: `--extra-device-vectorization` is deliberately omitted -- it causes
/// catastrophic register spill on the 3000-line experience kernel (sm_90),
/// blowing the per-thread stack and triggering CUDA_ERROR_ILLEGAL_ADDRESS.
const NVCC_OPT_FLAGS: &[&str] = &[
    "-O3",
    "--use_fast_math",
    "--ftz=true",
    "--fmad=true",
];

/// Compile CUDA source to a native cubin for the device behind `context`.
///
/// Uses a deterministic SHA-256 disk cache keyed on `(arch, source, flags)`.
/// First call for a given kernel + architecture compiles via `nvcc` (offline);
/// subsequent calls load the cached cubin in <10 ms.
///
/// The experience-collector kernel (3000+ lines of `.cu`) typically
/// takes ~15s to compile via nvcc on H100. With caching, subsequent runs load
/// in <10ms (file existence check only).
///
/// # Feature gate
/// Only available with the `cuda` feature.
///
/// # Compilation strategy
///
/// Uses nvcc (offline compiler) with full `-O3` optimization and `-cubin` output.
/// Produces native SASS for the exact GPU architecture (sm_XX), not virtual PTX.
/// Requires nvcc in `$CUDA_HOME/bin/` or `$PATH`.
pub fn compile_ptx_for_device(
    src: &str,
    context: &cudarc::driver::CudaContext,
) -> Result<cudarc::nvrtc::Ptx, String> {
    use cudarc::driver::sys::CUdevice_attribute;
    let major = context
        .attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR)
        .unwrap_or(0);
    let minor = context
        .attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR)
        .unwrap_or(0);

    // Build arch string for nvcc: "sm_XX" (real arch for native cubin output).
    // cubin is GPU-specific but eliminates driver JIT entirely.
    let arch_str: &'static str = match (major, minor) {
        (9, 0) => "sm_90",
        (9, _) => "sm_90",   // sm_90a and beyond
        (8, 9) => "sm_89",
        (8, 6) => "sm_86",
        (8, 0) => "sm_80",
        (7, 5) => "sm_75",
        (7, 0) => "sm_70",
        _ => "sm_80",        // safe default
    };

    // Try loading cached cubin first
    if let Some(cached) = load_cached_cubin(arch_str, src) {
        return Ok(cached);
    }

    // Cache miss -- compile via nvcc to native cubin
    let cubin = compile_cubin_with_nvcc(src, arch_str)?;

    Ok(cubin)
}

/// Compile CUDA source to native cubin using nvcc with maximum optimization.
///
/// See [`NVCC_OPT_FLAGS`] for the flags applied. `-cubin` and `-arch=sm_XX`
/// are added automatically. Output is a native SASS binary -- no driver JIT.
/// The cubin is written directly to the cache directory and loaded via
/// `Ptx::from_file()` which calls `cuModuleLoad()`.
fn compile_cubin_with_nvcc(
    src: &str,
    arch: &str,
) -> Result<cudarc::nvrtc::Ptx, String> {
    use std::io::Write;

    // Find nvcc: prefer $CUDA_HOME/bin/nvcc, fall back to PATH
    let nvcc = std::env::var("CUDA_HOME")
        .map(|home| std::path::PathBuf::from(home).join("bin/nvcc"))
        .unwrap_or_else(|_| std::path::PathBuf::from("nvcc"));

    // Write source to temp file (nvcc requires file input)
    let tmp_dir = std::env::temp_dir().join("foxhunt_nvcc");
    std::fs::create_dir_all(&tmp_dir)
        .map_err(|e| format!("Failed to create nvcc temp dir: {e}"))?;

    // Ensure cache dir exists -- nvcc writes cubin directly there
    let cache_dir = cubin_cache_dir();
    std::fs::create_dir_all(&cache_dir)
        .map_err(|e| format!("Failed to create cubin cache dir: {e}"))?;

    let key = cubin_cache_key(arch, src);
    let hash_prefix = key.get(..16).unwrap_or(&key);
    let src_path = tmp_dir.join(format!("{hash_prefix}.cu"));
    let cubin_path = cache_dir.join(format!("{key}.cubin"));

    {
        let mut f = std::fs::File::create(&src_path)
            .map_err(|e| format!("Failed to write CUDA source to {}: {e}", src_path.display()))?;
        f.write_all(src.as_bytes())
            .map_err(|e| format!("Failed to write CUDA source: {e}"))?;
    }

    // Invoke nvcc with -cubin (native SASS) instead of -ptx (virtual ISA).
    // This eliminates driver JIT entirely -- cuModuleLoad loads SASS directly.
    let arch_flag = format!("-arch={arch}");
    let out_path_str = cubin_path.to_str().unwrap_or("out.cubin");
    let in_path = src_path.to_str().unwrap_or("in.cu");
    let mut args: Vec<&str> = vec!["-cubin"];
    args.extend_from_slice(NVCC_OPT_FLAGS);
    args.extend_from_slice(&[&arch_flag, "-o", out_path_str, in_path]);
    let output = std::process::Command::new(&nvcc)
        .args(&args)
        .output()
        .map_err(|e| format!(
            "Failed to execute nvcc at {}: {e}. \
             Ensure CUDA toolkit is installed and nvcc is in $CUDA_HOME/bin/ or $PATH.",
            nvcc.display()
        ))?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        let stdout = String::from_utf8_lossy(&output.stdout);
        return Err(format!(
            "nvcc compilation failed (exit={}):\n{stderr}\n{stdout}",
            output.status.code().unwrap_or(-1)
        ));
    }

    // Verify the cubin was written
    let cubin_size = std::fs::metadata(&cubin_path)
        .map_err(|e| format!("Failed to read compiled cubin from {}: {e}", cubin_path.display()))?
        .len();

    // Clean up source temp file (best-effort), cubin stays in cache
    drop(std::fs::remove_file(&src_path));

    let flags_str = NVCC_OPT_FLAGS.join(" ");
    tracing::info!(
        arch,
        src_len = src.len(),
        cubin_len = cubin_size,
        flags = %flags_str,
        "nvcc cubin compilation succeeded (no driver JIT)"
    );

    // Load via Ptx::from_file -- cuModuleLoad() handles cubin natively
    Ok(cudarc::nvrtc::Ptx::from_file(&cubin_path))
}

/// Resolve the cubin cache directory.
///
/// Prefers `$CARGO_TARGET_DIR/.cubin_cache/` (persisted on CI PVC between runs).
/// Falls back to `/tmp/.cubin_cache/` when `CARGO_TARGET_DIR` is unset.
fn cubin_cache_dir() -> std::path::PathBuf {
    let base = std::env::var("CARGO_TARGET_DIR")
        .map(std::path::PathBuf::from)
        .unwrap_or_else(|_| std::path::PathBuf::from("/tmp"));
    base.join(".cubin_cache")
}

/// Compute deterministic SHA-256 cache key from (arch, source, flags).
///
/// Uses SHA-256 (not `DefaultHasher`) because `SipHash` uses random keys
/// per process -- the cache would never hit across separate cargo test runs.
fn cubin_cache_key(arch: &str, src: &str) -> String {
    use sha2::{Sha256, Digest};
    let mut h = Sha256::new();
    h.update(arch.as_bytes());
    h.update(b"\x00"); // separator
    // Include compile flags so any flag change auto-invalidates cached cubin.
    for flag in NVCC_OPT_FLAGS {
        h.update(flag.as_bytes());
        h.update(b"\x00");
    }
    h.update(src.as_bytes());
    format!("{:x}", h.finalize())
}

/// Try to load cached cubin from disk.
fn load_cached_cubin(
    arch: &str,
    src: &str,
) -> Option<cudarc::nvrtc::Ptx> {
    let key = cubin_cache_key(arch, src);
    let cache_path = cubin_cache_dir().join(format!("{key}.cubin"));
    if cache_path.exists() {
        let cubin_size = std::fs::metadata(&cache_path).map(|m| m.len()).unwrap_or(0);
        tracing::info!(
            "cubin cache HIT: {} ({} bytes)",
            cache_path.display(),
            cubin_size,
        );
        Some(cudarc::nvrtc::Ptx::from_file(&cache_path))
    } else {
        tracing::info!("cubin cache MISS: {}", cache_path.display());
        None
    }
}
