// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This software may be used and distributed according to the terms of the
// GNU General Public License version 2.

use anyhow::anyhow;
use anyhow::bail;
use anyhow::Context;
use anyhow::Result;
use glob::glob;
use libbpf_cargo::SkeletonBuilder;
use sscanf::sscanf;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::env;
use std::path::Path;
use std::path::PathBuf;
use std::process::Command;

lazy_static::lazy_static! {
    // Map clang archs to the __TARGET_ARCH list in
    // tools/lib/bpf/bpf_tracing.h in the kernel tree.
    static ref ARCH_MAP: HashMap<&'static str, &'static str> = vec![
    ("x86", "x86"),
    ("x86_64", "x86"),
    ("s390", "s390"),
    ("s390x", "s390"),
    ("arm", "arm"),
    ("aarch64", "arm64"),
    ("mips", "mips"),
    ("mips64", "mips"),
    ("ppc32", "powerpc"),
    ("ppc64", "powerpc"),
    ("ppc64le", "powerpc"),
    ("sparc", "sparc"),
    ("sparcv9", "sparc"),
    ("riscv32", "riscv"),
    ("riscv64", "riscv"),
    ("arc", "arc"),			// unsure this is supported
    ("loongarch64", "loongarch"),	// ditto
    ].into_iter().collect();
}

#[derive(Debug)]
/// # Build helpers for sched_ext schedulers with Rust userspace component
///
/// This is to be used from `build.rs` of a cargo project which implements a
/// [sched_ext](https://github.com/sched-ext/scx) scheduler with C BPF
/// component and Rust userspace component. `BpfBuilder` provides everything
/// necessary to build the BPF component and generate Rust bindings.
/// BpfBuilder provides the followings.
///
/// 1. *`vmlinux.h` and other common BPF header files*
///
/// All sched_ext BPF implementations require `vmlinux.h` and many make use
/// of common constructs such as
/// [`user_exit_info`](https://github.com/sched-ext/scx/blob/main/scheds/include/common/user_exit_info.h).
/// `BpfBuilder` makes these headers available when compiling BPF source
/// code and generating bindings for it. The included headers can be browsed
/// at <https://github.com/sched-ext/scx/tree/main/scheds/include>.
///
/// These headers can be superseded using environment variables which will
/// be discussed later.
///
/// 2. *Header bindings using `bindgen`*
///
/// If enabled with `.enable_intf()`, the input `.h` file is processed by
/// `bindgen` to generate Rust bindings. This is useful in establishing
/// shared constants and data types between the BPF and user components.
///
/// Note that the types generated with `bindgen` are different from the
/// types used by the BPF skeleton even when they are the same types in BPF.
/// This is a source of ugliness and we are hoping to address it by
/// improving `libbpf-cargo` in the future.
///
/// 3. *BPF compilation and generation of the skeleton and its bindings*
///
/// If enabled with `.enable_skel()`, the input `.bpf.c` file is compiled
/// and its skeleton and bindings are generated using `libbpf-cargo`.
///
/// ## An Example
///
/// This section shows how `BpfBuilder` can be used in an example project.
/// For a concrete example, take a look at
/// [`scx_rusty`](https://github.com/sched-ext/scx/tree/main/scheds/rust/scx_rusty).
///
/// A minimal source tree using all features would look like the following:
///
/// ```text
/// scx_hello_world
/// |-- Cargo.toml
/// |-- build.rs
/// \-- src
///     |-- main.rs
///     |-- bpf_intf.rs
///     |-- bpf_skel.rs
///     \-- bpf
///         |-- intf.h
///         \-- main.c
/// ```
///
/// The following three files would contain the actual implementation:
///
/// - `src/main.rs`: Rust userspace component which loads the BPF blob and
/// interacts it using the generated bindings.
///
/// - `src/bpf/intf.h`: C header file definining constants and structs
/// that will be used by both the BPF and userspace components.
///
/// - `src/bpf/main.c`: C source code implementing the BPF component -
/// including `struct sched_ext_ops`.
///
/// And then there are boilerplates to generate the bindings and make them
/// available as modules to `main.rs`.
///
/// - `Cargo.toml`: Includes `scx_utils` in the `[build-dependencies]`
/// section.
///
/// - `build.rs`: Uses `scx_utils::BpfBuilder` to build and generate
/// bindings for the BPF component. For this project, it can look like the
/// following.
///
/// ```should_panic
/// fn main() {
///     scx_utils::BpfBuilder::new()
///         .unwrap()
///         .enable_intf("src/bpf/intf.h", "bpf_intf.rs")
///         .enable_skel("src/bpf/main.bpf.c", "bpf")
///         .build()
///         .unwrap();
/// }
/// ```
///
/// - `bpf_intf.rs`: Import the bindings generated by `bindgen` into a
/// module. Above, we told `.enable_intf()` to generate the bindings into
/// `bpf_intf.rs`, so the file would look like the following. The `allow`
/// directives are useful if the header is including `vmlinux.h`.
///
/// ```ignore
/// #![allow(non_upper_case_globals)]
/// #![allow(non_camel_case_types)]
/// #![allow(non_snake_case)]
/// #![allow(dead_code)]
///
/// include!(concat!(env!("OUT_DIR"), "/bpf_intf.rs"));
/// ```
///
/// - `bpf_skel.rs`: Import the BPF skeleton bindings generated by
/// `libbpf-cargo` into a module. Above, we told `.enable_skel()` to use the
/// skeleton name `bpf`, so the file would look like the following.
///
/// ```ignore
/// include!(concat!(env!("OUT_DIR"), "/bpf_skel.rs"));
/// ```
///
/// ## Compiler Flags and Environment Variables
///
/// BPF being its own CPU architecture and independent runtime environment,
/// build environment and steps are already rather complex. The need to
/// interface between two different languages - C and Rust - adds further
/// complexities. `BpfBuilder` automates most of the process. The determined
/// build environment is recorded in the `build.rs` output and can be
/// obtained with a command like the following:
///
/// ```text
/// $ grep '^scx_utils:clang=' target/release/build/scx_rusty-*/output
/// ```
///
/// While the automatic settings should work most of the time, there can be
/// times when overriding them is necessary. The following environment
/// variables can be used to customize the build environment.
///
/// - `BPF_CLANG`: The clang command to use. (Default: `clang`)
///
/// - `BPF_CFLAGS`: Compiler flags to use when building BPF source code. If
///   specified, the flags from this variable are the only flags passed to
///   the compiler. `BpfBuilder` won't generate any flags including `-I`
///   flags for the common header files and other `CFLAGS` related variables
///   are ignored.
///
/// - `BPF_BASE_CFLAGS`: Override the non-include part of cflags.
///
/// - `BPF_EXTRA_CFLAGS_PRE_INCL`: Add cflags before the automic include
///   search path options. Header files in the search paths added by this
///   variable will supercede the automatic ones.
///
/// - `BPF_EXTRA_CFLAGS_POST_INCL`: Add cflags after the automic include
///   search path options. Header paths added by this variable will be
///   searched only if the target header file can't be found in the
///   automatic header paths.
///
/// - `RUSTFLAGS`: This is a generic `cargo` flag and can be useful for
///   specifying extra linker flags.
///
/// A common case for using the above flags is using the latest `libbpf`
/// from the kernel tree. Let's say the kernel tree is at `$KERNEL` and
/// `libbpf`. The following builds `libbpf` shipped with the kernel:
///
/// ```test
/// $ cd $KERNEL
/// $ make -C tools/bpf/bpftool
/// ```
///
/// To link the scheduler against the resulting `libbpf`:
///
/// ```test
/// $ env BPF_EXTRA_CFLAGS_POST_INCL=$KERNEL/tools/bpf/bpftool/libbpf/include \
///   RUSTFLAGS="-C link-args=-lelf -C link-args=-lz -C link-args=-lzstd \
///   -L$KERNEL/tools/bpf/bpftool/libbpf" cargo build --release
/// ```
pub struct BpfBuilder {
    clang: (String, String, String), // (clang, ver, arch)
    cflags: Vec<String>,
    out_dir: PathBuf,

    intf_input_output: Option<(String, String)>,
    skel_input_name: Option<(String, String)>,
    skel_deps: Option<Vec<String>>,
}

impl BpfBuilder {
    fn skip_clang_version_prefix(line: &str) -> &str {
        if let Some(index) = line.find("clang version") {
            &line[index..]
        } else {
            line
        }
    }

    fn find_clang() -> Result<(String, String, String)> {
        let clang = env::var("BPF_CLANG").unwrap_or("clang".into());
        let output = Command::new(&clang)
            .args(["--version"])
            .output()
            .with_context(|| format!("Failed to run \"{} --version\"", &clang))?;

        let stdout = String::from_utf8(output.stdout)?;
        let (mut ver, mut arch) = (None, None);
        for line in stdout.lines() {
            if let Ok(v) = sscanf!(
                Self::skip_clang_version_prefix(line),
                "clang version {String}"
            ) {
                // Version could be followed by (URL SHA1). Only take
                // the first word.
                ver = Some(v.split_whitespace().next().unwrap().to_string());
                continue;
            }
            if let Ok(v) = sscanf!(line, "Target: {String}") {
                arch = Some(v.split('-').next().unwrap().to_string());
                continue;
            }
        }

        let (ver, arch) = (
            ver.ok_or(anyhow!("Failed to read clang version"))?,
            arch.ok_or(anyhow!("Failed to read clang target arch"))?,
        );

        if version_compare::compare(&ver, "16") == Ok(version_compare::Cmp::Lt) {
            bail!(
                "clang < 16 loses high 32 bits of 64 bit enums when compiling BPF ({:?} ver={:?})",
                &clang,
                &ver
            );
        }
        if version_compare::compare(&ver, "17") == Ok(version_compare::Cmp::Lt) {
            println!(
                "cargo:warning=clang >= 17 recommended ({:?} ver={:?})",
                &clang, &ver
            );
        }

        Ok((clang, ver, arch))
    }

    fn determine_base_cflags(
        (clang, _ver, arch): &(String, String, String),
    ) -> Result<Vec<String>> {
        // Determine kernel target arch.
        let kernel_target = match ARCH_MAP.get(arch.as_str()) {
            Some(v) => v,
            None => bail!("CPU arch {:?} not found in ARCH_MAP", &arch),
        };

        // Determine system includes.
        let output = Command::new(&clang)
            .args(["-v", "-E", "-"])
            .output()
            .with_context(|| format!("Failed to run \"{} -v -E - < /dev/null", &clang))?;
        let stderr = String::from_utf8(output.stderr)?;

        let mut sys_incls = None;
        for line in stderr.lines() {
            if line == "#include <...> search starts here:" {
                sys_incls = Some(vec![]);
                continue;
            }
            if sys_incls.is_none() {
                continue;
            }
            if line == "End of search list." {
                break;
            }

            sys_incls.as_mut().unwrap().push(line.trim());
        }
        let sys_incls = match sys_incls {
            Some(v) => v,
            None => bail!("Failed to find system includes from {:?}", &clang),
        };

        // Determine endian.
        let output = Command::new(&clang)
            .args(["-dM", "-E", "-"])
            .output()
            .with_context(|| format!("Failed to run \"{} -dM E - < /dev/null", &clang))?;
        let stdout = String::from_utf8(output.stdout)?;

        let mut endian = None;
        for line in stdout.lines() {
            match sscanf!(line, "#define __BYTE_ORDER__ {str}") {
                Ok(v) => {
                    endian = Some(match v {
                        "__ORDER_LITTLE_ENDIAN__" => "little",
                        "__ORDER_BIG_ENDIAN__" => "big",
                        v => bail!("Unknown __BYTE_ORDER__ {:?}", v),
                    });
                    break;
                }
                _ => {}
            }
        }
        let endian = match endian {
            Some(v) => v,
            None => bail!("Failed to find __BYTE_ORDER__ from {:?}", &clang),
        };

        // Assemble cflags.
        let mut cflags: Vec<String> =
            ["-g", "-O2", "-Wall", "-Wno-compare-distinct-pointer-types'"]
                .into_iter()
                .map(|x| x.into())
                .collect();
        cflags.push(format!("-D__TARGET_ARCH_{}", &kernel_target));
        cflags.push("-mcpu=v3".into());
        cflags.push(format!("-m{}-endian", endian));
        cflags.append(
            &mut sys_incls
                .into_iter()
                .flat_map(|x| ["-idirafter".into(), x.into()])
                .collect(),
        );
        Ok(cflags)
    }

    const BPF_H_TAR: &'static [u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bpf_h.tar"));

    fn install_bpf_h<P: AsRef<Path>>(dest: P) -> Result<()> {
        let mut ar = tar::Archive::new(Self::BPF_H_TAR);
        ar.unpack(dest)?;
        Ok(())
    }

    /// Return `(VER, SHA1)` from which the bulit-in `vmlinux.h` is generated.
    pub fn vmlinux_h_ver_sha1() -> (String, String) {
        let mut ar = tar::Archive::new(Self::BPF_H_TAR);

        for file in ar.entries().unwrap() {
            let file = file.unwrap();
            if file.header().path().unwrap() != Path::new("vmlinux/vmlinux.h") {
                continue;
            }

            let name = file
                .link_name()
                .unwrap()
                .unwrap()
                .to_string_lossy()
                .to_string();

            return sscanf!(name, "vmlinux-v{String}-g{String}.h").unwrap();
        }

        panic!("vmlinux/vmlinux.h not found");
    }

    fn determine_cflags<P>(clang: &(String, String, String), out_dir: P) -> Result<Vec<String>>
    where
        P: AsRef<Path> + std::fmt::Debug,
    {
        let bpf_h = out_dir
            .as_ref()
            .join("scx_utils-bpf_h")
            .to_str()
            .ok_or(anyhow!(
                "{:?}/scx_utils-bph_h can't be converted to str",
                &out_dir
            ))?
            .to_string();
        Self::install_bpf_h(&bpf_h)?;

        let mut cflags = Vec::<String>::new();

        cflags.append(&mut match env::var("BPF_BASE_CFLAGS") {
            Ok(v) => v.split_whitespace().map(|x| x.into()).collect(),
            _ => Self::determine_base_cflags(&clang)?,
        });

        cflags.append(&mut match env::var("BPF_EXTRA_CFLAGS_PRE_INCL") {
            Ok(v) => v.split_whitespace().map(|x| x.into()).collect(),
            _ => vec![],
        });

        cflags.push(format!("-I{}", &bpf_h));
        cflags.push(format!("-I{}/vmlinux", &bpf_h));
        cflags.push(format!("-I{}/bpf-compat", &bpf_h));

        cflags.append(&mut match env::var("BPF_EXTRA_CFLAGS_POST_INCL") {
            Ok(v) => v.split_whitespace().map(|x| x.into()).collect(),
            _ => vec![],
        });

        Ok(cflags)
    }

    /// Create a new `BpfBuilder` struct. Call `enable` and `set` methods to
    /// configure and `build` method to compile and generate bindings. See
    /// the struct documentation for details.
    pub fn new() -> Result<Self> {
        let out_dir = PathBuf::from(env::var("OUT_DIR")?);

        let clang = Self::find_clang()?;
        let cflags = match env::var("BPF_CFLAGS") {
            Ok(v) => v.split_whitespace().map(|x| x.into()).collect(),
            _ => Self::determine_cflags(&clang, &out_dir)?,
        };

        println!("scx_utils:clang={:?} {:?}", &clang, &cflags);

        Ok(Self {
            clang,
            cflags,
            out_dir,

            intf_input_output: None,
            skel_input_name: None,
            skel_deps: None,
        })
    }

    /// Enable generation of header bindings using `bindgen`. `@input` is
    /// the `.h` file defining the constants and types to be shared between
    /// BPF and Rust components. `@output` is the `.rs` file to be
    /// generated.
    pub fn enable_intf(&mut self, input: &str, output: &str) -> &mut Self {
        self.intf_input_output = Some((input.into(), output.into()));
        self
    }

    /// Enable compilation of BPF code and generation of the skeleton and
    /// its Rust bindings. `@input` is the `.bpf.c` file containing the BPF
    /// source code and `@output` is the `.rs` file to be generated.
    pub fn enable_skel(&mut self, input: &str, name: &str) -> &mut Self {
        self.skel_input_name = Some((input.into(), name.into()));
        self
    }

    /// By default, all `.[hc]` files in the same directory as the source
    /// BPF `.c` file are treated as dependencies and the skeleton is
    /// regenerated if any has changed. This method replaces the automatic
    /// dependencies with `@deps`.
    pub fn set_skel_deps<'a, I>(&mut self, deps: I) -> &mut Self
    where
        I: IntoIterator<Item = &'a str>,
    {
        self.skel_deps = Some(deps.into_iter().map(|d| d.to_string()).collect());
        self
    }

    fn bindgen_bpf_intf(&self, deps: &mut BTreeSet<String>) -> Result<()> {
        let (input, output) = match &self.intf_input_output {
            Some(pair) => pair,
            None => return Ok(()),
        };

        // Tell cargo to invalidate the built crate whenever the wrapper changes
        deps.insert(input.to_string());

	// FIXME - bindgen's API changed between 0.68 and 0.69 so that
	// `bindgen::CargoCallbacks::new()` should be used instead of
	// `bindgen::CargoCallbacks`. Unfortunately, as of Dec 2023, fedora
	// is shipping 0.68. To accommodate fedora, allow both 0.68 and 0.69
	// of bindgen and suppress deprecation warning. Remove the following
	// once fedora can be updated to bindgen >= 0.69.
	#[allow(deprecated)]

        // The bindgen::Builder is the main entry point to bindgen, and lets
        // you build up options for the resulting bindings.
        let bindings = bindgen::Builder::default()
            // Should run clang with the same -I options as BPF compilation.
            .clang_args(
                self.cflags
                    .iter()
                    .chain(["-target".into(), "bpf".into()].iter()),
            )
            // The input header we would like to generate bindings for.
            .header(input)
            // Tell cargo to invalidate the built crate whenever any of the
            // included header files changed.
            .parse_callbacks(Box::new(bindgen::CargoCallbacks))
            .generate()
            .context("Unable to generate bindings")?;

        bindings
            .write_to_file(self.out_dir.join(output))
            .context("Couldn't write bindings")
    }

    fn gen_bpf_skel(&self, deps: &mut BTreeSet<String>) -> Result<()> {
        let (input, name) = match &self.skel_input_name {
            Some(pair) => pair,
            None => return Ok(()),
        };

        let obj = self.out_dir.join(format!("{}.bpf.o", name));
        let skel_path = self.out_dir.join(format!("{}_skel.rs", name));

        SkeletonBuilder::new()
            .source(input)
            .obj(&obj)
            .clang(&self.clang.0)
            .clang_args(&self.cflags)
            .build_and_generate(&skel_path)?;

        match &self.skel_deps {
            Some(skel_deps) => {
                for path in skel_deps {
                    deps.insert(path.to_string());
                }
            }
            None => {
                let c_path = PathBuf::from(input);
                let dir = c_path
                    .parent()
                    .ok_or(anyhow!("Source {:?} doesn't have parent dir", c_path))?
                    .to_str()
                    .ok_or(anyhow!("Parent dir of {:?} isn't a UTF-8 string", c_path))?;

                for path in glob(&format!("{}/*.[hc]", dir))?.filter_map(Result::ok) {
                    deps.insert(
                        path.to_str()
                            .ok_or(anyhow!("Path {:?} is not a valid string", path))?
                            .to_string(),
                    );
                }
            }
        }
        Ok(())
    }

    /// Build and generate the enabled bindings.
    pub fn build(&self) -> Result<()> {
        let mut deps = BTreeSet::new();

        self.bindgen_bpf_intf(&mut deps)?;
        self.gen_bpf_skel(&mut deps)?;

        println!("cargo:rerun-if-env-changed=BPF_CLANG");
        println!("cargo:rerun-if-env-changed=BPF_CFLAGS");
        println!("cargo:rerun-if-env-changed=BPF_BASE_CFLAGS");
        println!("cargo:rerun-if-env-changed=BPF_EXTRA_CFLAGS_PRE_INCL");
        println!("cargo:rerun-if-env-changed=BPF_EXTRA_CFLAGS_POST_INCL");
        for dep in deps.iter() {
            println!("cargo:rerun-if-changed={}", dep);
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn test_bpf_builder_new() {
        let res = super::BpfBuilder::new();
        assert!(res.is_ok(), "Failed to create BpfBuilder ({:?})", &res);
    }

    #[test]
    fn test_vmlinux_h_ver_sha1() {
        let (ver, sha1) = super::BpfBuilder::vmlinux_h_ver_sha1();

        println!("vmlinux.h: ver={:?} sha1={:?}", &ver, &sha1,);

        assert!(
            regex::Regex::new(r"^[1-9][0-9]*\.[1-9][0-9]*(\.[1-9][0-9]*)?$")
                .unwrap()
                .is_match(&ver)
        );
        assert!(regex::Regex::new(r"^[0-9a-z]{12}$")
            .unwrap()
            .is_match(&sha1));
    }
}