//! Build script for Flash Attention CUDA kernels
//! Optimized for RTX 5090 (sm_89/sm_90) architecture

use std::env;
use std::path::{Path, PathBuf};
use std::process::Command;

fn main() {
    println!("cargo:rerun-if-changed=cuda/");
    println!("cargo:rerun-if-changed=metal/");
    println!("cargo:rerun-if-changed=build.rs");

    // Only compile CUDA if the cuda feature is enabled
    let cuda_feature = env::var("CARGO_FEATURE_CUDA").is_ok();
    let metal_feature = env::var("CARGO_FEATURE_METAL").is_ok();

    if cuda_feature {
        let cuda_available = detect_cuda();
        if cuda_available {
            println!("cargo:rustc-cfg=cuda_available");
            compile_cuda_kernels();
            link_cuda_libraries();
        } else {
            println!("cargo:warning=CUDA feature enabled but CUDA not found on system");
        }
    }

    if metal_feature {
        println!("cargo:rustc-cfg=metal_available");
        compile_metal_shaders();
        println!("cargo:warning=Metal GPU support enabled for macOS");
    }

    if !cuda_feature && !metal_feature {
        println!(
            "cargo:warning=No GPU backend enabled. Enable 'cuda' or 'metal' feature for GPU acceleration"
        );
    }
}

fn detect_cuda() -> bool {
    // Try to find nvcc compiler
    if let Ok(output) = Command::new("nvcc").arg("--version").output()
        && output.status.success()
    {
        let version_output = String::from_utf8_lossy(&output.stdout);
        println!(
            "cargo:warning=Found CUDA: {}",
            version_output.lines().next().unwrap_or("unknown")
        );
        return true;
    }

    // Check for CUDA_PATH environment variable
    if let Ok(cuda_path) = env::var("CUDA_PATH") {
        let nvcc_path = Path::new(&cuda_path).join("bin").join("nvcc");
        if nvcc_path.exists() {
            println!("cargo:warning=Found CUDA at: {cuda_path}");
            return true;
        }
    }

    false
}

fn compile_cuda_kernels() {
    let out_dir = env::var("OUT_DIR").unwrap();
    let cuda_dir = PathBuf::from("cuda");

    // RTX 5090 optimized compilation flags
    let cuda_flags = vec![
        "--gpu-architecture=compute_90", // RTX 5090 architecture (Ada Lovelace)
        "--gpu-code=sm_90",              // RTX 5090 compute capability
        "--use_fast_math",
        "--maxrregcount=255",
        "-O3",
        "--extra-device-vectorization",
        "--fmad=true",
        "--prec-div=false",
        "--prec-sqrt=false",
        "--ftz=true", // Flush denormal to zero for performance
        "-DRTX_5090_OPTIMIZATION=1",
        "-DUSE_TENSOR_CORE_4TH_GEN=1",
        "-DMAX_SHARED_MEMORY=163840", // 160KB for RTX 5090
    ];

    let kernels = vec![
        "flash_attention_forward.cu",
        "flash_attention_backward.cu",
        "online_softmax.cu",
        "utils.cu",
    ];

    for kernel in kernels {
        let input_path = cuda_dir.join(kernel);
        let output_name = kernel.replace(".cu", ".ptx");
        let output_path = Path::new(&out_dir).join(&output_name);

        println!(
            "cargo:warning=Compiling CUDA kernel: {} -> {}",
            input_path.display(),
            output_path.display()
        );

        let mut cmd = Command::new("nvcc");
        cmd.arg("--ptx")
            .arg(&input_path)
            .arg("-o")
            .arg(&output_path);

        // Add optimization flags
        for flag in &cuda_flags {
            cmd.arg(flag);
        }

        // Add include directories
        if let Ok(cuda_path) = env::var("CUDA_PATH") {
            cmd.arg(format!("-I{cuda_path}/include"));
        }

        let output = cmd.output().expect("Failed to execute nvcc");

        if output.status.success() {
            println!("cargo:warning=Successfully compiled {kernel}");
        } else {
            let stderr = String::from_utf8_lossy(&output.stderr);
            panic!("Failed to compile CUDA kernel {kernel}: {stderr}");
        }
    }

    // Create a combined PTX file with all kernels
    combine_ptx_files(&out_dir);
}

fn combine_ptx_files(out_dir: &str) {
    use std::fs::{File, OpenOptions};
    use std::io::{BufWriter, Read, Write};

    let combined_path = Path::new(out_dir).join("flash_attention_kernels.ptx");
    let mut combined_file = BufWriter::new(
        OpenOptions::new()
            .create(true)
            .write(true)
            .truncate(true)
            .open(&combined_path)
            .expect("Failed to create combined PTX file"),
    );

    let ptx_files = vec![
        "flash_attention_forward.ptx",
        "flash_attention_backward.ptx",
        "online_softmax.ptx",
        "utils.ptx",
    ];

    // Write PTX header
    writeln!(combined_file, ".version 7.8").unwrap();
    writeln!(combined_file, ".target sm_90").unwrap();
    writeln!(combined_file, ".address_size 64").unwrap();
    writeln!(combined_file).unwrap();

    for ptx_file in ptx_files {
        let ptx_path = Path::new(out_dir).join(ptx_file);
        if ptx_path.exists() {
            let mut contents = String::new();
            File::open(&ptx_path)
                .expect("Failed to open PTX file")
                .read_to_string(&mut contents)
                .expect("Failed to read PTX file");

            // Skip header lines and write content
            let lines: Vec<&str> = contents.lines().collect();
            let content_start = lines.iter().position(|&line| line.is_empty()).unwrap_or(0) + 1;

            for line in &lines[content_start..] {
                writeln!(combined_file, "{line}").unwrap();
            }

            writeln!(combined_file).unwrap();
        }
    }

    combined_file.flush().unwrap();
    drop(combined_file);

    println!(
        "cargo:warning=Created combined PTX file: {}",
        combined_path.display()
    );

    // Tell Rust to include the PTX file
    println!(
        "cargo:rustc-env=FLASH_ATTENTION_PTX_PATH={}",
        combined_path.display()
    );
}

fn link_cuda_libraries() {
    // Add CUDA library paths
    if let Ok(cuda_path) = env::var("CUDA_PATH") {
        println!("cargo:rustc-link-search=native={cuda_path}/lib64");
    } else {
        // Try common CUDA installation paths
        let cuda_paths = vec![
            "/usr/local/cuda/lib64",
            "/opt/cuda/lib64",
            "/usr/lib/x86_64-linux-gnu",
        ];

        for path in cuda_paths {
            if Path::new(path).exists() {
                println!("cargo:rustc-link-search=native={path}");
                break;
            }
        }
    }

    // Link CUDA libraries
    println!("cargo:rustc-link-lib=cudart");
    println!("cargo:rustc-link-lib=cublas");
    println!("cargo:rustc-link-lib=curand");
    println!("cargo:rustc-link-lib=cusparse");
    println!("cargo:rustc-link-lib=cufft");

    // RTX 5090 specific optimizations
    println!("cargo:rustc-cfg=rtx_5090_optimized");
    println!("cargo:rustc-cfg=tensor_core_4th_gen");
}

fn compile_metal_shaders() {
    // Metal shaders can be compiled using metallib
    // For now, we'll just mark them as available
    // In a full implementation, you would compile .metal files to .metallib

    let _out_dir = env::var("OUT_DIR").unwrap();
    let metal_dir = PathBuf::from("metal");

    // Check if Metal shader source exists
    let flash_attention_metal = metal_dir.join("flash_attention.metal");
    if flash_attention_metal.exists() {
        println!(
            "cargo:warning=Found Metal shader source: {}",
            flash_attention_metal.display()
        );

        // In production, compile with:
        // xcrun -sdk macosx metal -c flash_attention.metal -o flash_attention.air
        // xcrun -sdk macosx metallib flash_attention.air -o flash_attention.metallib

        // For now, just copy the source path
        let metal_source_path = std::fs::canonicalize(&flash_attention_metal)
            .expect("Failed to get absolute path to Metal shader");
        println!(
            "cargo:rustc-env=METAL_SHADER_PATH={}",
            metal_source_path.display()
        );
        println!(
            "cargo:warning=Metal shader source available at: {}",
            metal_source_path.display()
        );
    } else {
        println!("cargo:warning=Metal shader source not found, using stub implementation");
    }
}