prebindgen 0.5.1

use std::{
    cell::RefCell,
    collections::{HashMap, HashSet},
    fs,
    path::{Path, PathBuf},
};

use itertools::Itertools;
use roxygen::roxygen;

use crate::{
    api::{batching::cfg_filter, record::Record, utils::jsonl::read_jsonl_file},
    SourceLocation, TargetTriple, CRATE_NAME_FILE, FEATURES_FILE,
};

/// File extension for data files
const JSONL_EXTENSION: &str = ".jsonl";
thread_local! {
    static DOCTEST_SOURCE: RefCell<Option<Source>> = const { RefCell::new(None) };
}

/// Provides access to prebindgen data generated by the `#[prebindgen]` macro
///
/// The `Source` struct reads prebindgen data files from a directory that was
/// initialized by [`init_prebindgen_out_dir`](crate::init_prebindgen_out_dir) in
/// the source project's `build.rs`. This allows destination projects to access
/// and process the collected FFI interface definitions.
///
/// # Workflow
///
/// 1. **Source project**: Calls `init_prebindgen_out_dir()` in `build.rs` and uses `#[prebindgen]` macro
/// 2. **Destination project**: Uses `Source::new()` to read the collected data
///
/// # Groups
///
/// Items are organized into groups based on the first parameter of the `#[prebindgen]`
/// macro. For example, `#[prebindgen("structs")]` creates items in the "structs" group.
/// Items without an explicit group are placed in the "default" group.
///
/// # Example
///
/// The `PREBINDGEN_OUT_DIR` constant is defined in the source FFI crate using the
/// `prebindgen_out_dir!()` macro:
///
/// ```rust
/// // In source_ffi/src/lib.rs
/// use prebindgen_proc_macro::{features, prebindgen, prebindgen_out_dir};
///
/// pub const PREBINDGEN_OUT_DIR: &str = prebindgen_out_dir!();
/// pub const FEATURES: &str = features!();
///
/// #[prebindgen]
/// pub fn my_function() -> i32 { 42 }
/// ```
///
/// Then in the destination project's `build.rs`:
///
/// ```
/// # prebindgen::Source::init_doctest_simulate();
/// let source = prebindgen::Source::new("source_ffi");
///
/// // Process all items
/// for (item, location) in source.items_all() {
///     // Process each syn::Item...
/// }
/// ```
#[derive(Clone)]
pub struct Source {
    crate_name: String,
    items: HashMap<String, Vec<(syn::Item, SourceLocation)>>,
    // Configuration needed to build a CfgFilter at iteration time
    features_constant: Option<String>,
    target_triple: Option<String>,
    features_list: Vec<String>, // normalized list from features.txt
}

impl Source {
    /// Create a new `Source` instance from specified directory
    #[roxygen]
    pub fn new<P: AsRef<Path>>(
        /// Path to the directory containing prebindgen data files
        input_dir: P,
    ) -> Self {
        // Shorthand for `builder(...).build()`: defaults to asserting the `FEATURES` constant
        Self::builder(input_dir).build()
    }

    /// Create a builder to configure how `Source` is constructed
    pub fn builder<P: AsRef<Path>>(input_dir: P) -> Builder {
        Builder::new(input_dir)
    }
}

impl Source {
    fn build_internal(
        input_dir: &Path,
        features_constant: Option<String>,
        target_triple: Option<String>,
        crate_name_override: Option<String>,
    ) -> Self {
        if let Some(source) = DOCTEST_SOURCE.with(|source| (*source.borrow()).clone()) {
            return source;
        }
        if !input_dir.is_dir() {
            panic!(
                "Input directory {} does not exist or is not a directory",
                input_dir.display()
            );
        }
        // The stored name (CARGO_PKG_NAME at capture time) doubles as the
        // "directory was initialized" check, so it is read even when
        // overridden. The override wins: it is the name THIS crate
        // references the source crate by (a Cargo.toml dependency rename).
        let stored_crate_name = read_stored_crate_name(input_dir).unwrap_or_else(|| {
            panic!(
                "The directory {} was not initialized with init_prebindgen_out_dir(). \
                Please ensure that init_prebindgen_out_dir() is called in the build.rs of the source crate.",
                input_dir.display()
            )
        });
        let crate_name = crate_name_override.unwrap_or(stored_crate_name);

        let groups = Self::discover_groups(input_dir);
        let mut items = HashMap::new();
        for group in groups {
            let records = Self::read_group(input_dir, &group);
            let group_items = records
                .iter()
                .map(|r| {
                    // Stamp the origin crate into every item's location: the
                    // captured JSONL doesn't carry it (the proc-macro runs
                    // inside the crate), but from here on the item streams
                    // are self-describing — streams from several sources can
                    // be chained into one `Flat::builder().items(..)` call
                    // without losing per-item origins.
                    let (item, mut loc) = r.parse();
                    loc.crate_name = Some(crate_name.clone());
                    (item, loc)
                })
                .collect::<Vec<_>>();
            items.insert(group, group_items);
        }

        // Read features list once and store normalized list
        let features_list = read_features_from_out_dir(input_dir);

        Self {
            crate_name,
            items,
            features_constant,
            features_list,
            target_triple,
        }
    }

    #[doc(hidden)]
    /// Initialize the Source for doctests by creating a dummy Source with a few test items.
    pub fn init_doctest_simulate() {
        let source = Self {
            crate_name: "source_ffi".to_string(),
            items: HashMap::from([
                (
                    "structs".to_string(),
                    vec![(
                        syn::parse_quote! {
                            #[prebindgen("structs")]
                            pub struct TestStruct {
                                pub field: i32,
                            }
                        },
                        SourceLocation {
                            crate_name: Some("source_ffi".to_string()),
                            ..SourceLocation::default()
                        },
                    )],
                ),
                (
                    "functions".to_string(),
                    vec![(
                        syn::parse_quote! {
                            #[prebindgen("functions")]
                            pub fn test_function() -> i32 { 42 }
                        },
                        SourceLocation {
                            crate_name: Some("source_ffi".to_string()),
                            ..SourceLocation::default()
                        },
                    )],
                ),
            ]),
            features_constant: None,
            target_triple: None,
            features_list: Vec::new(),
        };
        DOCTEST_SOURCE.with(|cell| {
            *cell.borrow_mut() = Some(source);
        });
    }

    /// Returns the name of the source crate that generated the prebindgen data
    ///
    /// # Example
    ///
    /// ```
    /// # prebindgen::Source::init_doctest_simulate();
    /// let source = prebindgen::Source::new("source_ffi");
    /// let crate_name = source.crate_name();
    /// ```
    pub fn crate_name(&self) -> &str {
        &self.crate_name
    }

    /// Returns an iterator over items from specific groups
    ///
    /// # Parameters
    ///
    /// * `groups` - Array of group names to include
    ///
    /// # Example
    ///
    /// ```
    /// # prebindgen::Source::init_doctest_simulate();
    /// let source = prebindgen::Source::new("source_ffi");
    /// // Process only items from "structs" and "functions" groups
    /// let items = source.items_in_groups(&["structs"]).collect::<Vec<_>>();
    /// assert_eq!(items.len(), 1); // only TestStruct should be present
    /// ```
    pub fn items_in_groups<'a>(
        &'a self,
        groups: &'a [&'a str],
    ) -> impl Iterator<Item = (syn::Item, SourceLocation)> + 'a {
        // Build a cfg filter and apply it lazily with itertools::batching
        let mut filter = self.build_cfg_filter();
        groups
            .iter()
            .filter_map(|group| self.items.get(*group))
            .flat_map(|records| records.iter().cloned())
            .batching(move |iter| filter.call(iter))
    }

    /// Returns an iterator over items excluding specific groups
    ///
    /// # Parameters
    ///
    /// * `groups` - Array of group names to exclude
    ///
    /// # Example
    ///
    /// ```
    /// # prebindgen::Source::init_doctest_simulate();
    /// let source = prebindgen::Source::new("source_ffi");
    /// let items = source.items_except_groups(&["structs"]).collect::<Vec<_>>();
    /// assert_eq!(items.len(), 1); // only test_function should be present
    /// ```
    pub fn items_except_groups<'a>(
        &'a self,
        groups: &'a [&'a str],
    ) -> impl Iterator<Item = (syn::Item, SourceLocation)> + 'a {
        // Build a cfg filter and apply it lazily with itertools::batching
        let mut filter = self.build_cfg_filter();
        self.items
            .iter()
            .filter(|(group, _)| !groups.contains(&group.as_str()))
            .flat_map(|(_, records)| records.iter().cloned())
            .batching(move |iter| filter.call(iter))
    }

    /// Returns an iterator over all items from all groups
    ///
    /// This is the most commonly used method for processing all prebindgen items.
    ///
    /// # Example
    ///
    /// ```
    /// # prebindgen::Source::init_doctest_simulate();
    /// let source = prebindgen::Source::new("source_ffi");
    /// let items: Vec<_> = source.items_all().collect();
    /// assert_eq!(items.len(), 2); // should contain TestStruct and test_function
    /// ```
    pub fn items_all(&self) -> impl Iterator<Item = (syn::Item, SourceLocation)> + '_ {
        // Build a cfg filter and apply it lazily with itertools::batching
        let mut filter = self.build_cfg_filter();
        self.items
            .values()
            .flat_map(|records| records.iter().cloned())
            .batching(move |iter| filter.call(iter))
    }

    /// Internal: construct a CfgFilter from the stored configuration and features file
    fn build_cfg_filter(&self) -> cfg_filter::CfgFilter {
        let mut builder = cfg_filter::CfgFilter::builder();
        if let Some(const_name) = &self.features_constant {
            // If the provided constant isn't fully qualified, qualify it with the crate name
            let qualified_const = if const_name.contains("::") {
                const_name.clone()
            } else {
                format!("{}::{}", self.crate_name.replace('-', "_"), const_name)
            };
            // Prepend each feature with "crate_name/" and join with spaces
            let features_list = self
                .features_list
                .iter()
                .map(|f| format!("{}/{}", self.crate_name, f))
                .join(" ");
            builder = builder.predefined_features(qualified_const, features_list);
        }
        if let Some(target) = &self.target_triple {
            let target_triple = TargetTriple::parse(target)
                .unwrap_or_else(|e| panic!("Failed to parse target triple '{}': {}", target, e));
            if let Some(arch) = target_triple.arch() {
                builder = builder.enable_target_arch(arch);
            }
            if let Some(vendor) = target_triple.vendor() {
                builder = builder.enable_target_vendor(vendor);
            }
            if let Some(os) = target_triple.os() {
                builder = builder.enable_target_os(os);
            }
            if let Some(env) = target_triple.env() {
                builder = builder.enable_target_env(env);
            }
        }
        builder.build()
    }

    /// Internal method to read all exported files matching the group name pattern `<group>_*`
    fn read_group<P: AsRef<Path>>(input_dir: P, group: &str) -> Vec<Record> {
        let pattern = format!("{group}_");
        let mut record_map = HashMap::new();

        // Read the directory and find all matching files
        if let Ok(entries) = fs::read_dir(&input_dir) {
            for entry in entries.flatten() {
                let path = entry.path();
                if let Some(file_name) = path.file_name().and_then(|n| n.to_str()) {
                    if file_name.starts_with(&pattern) && file_name.ends_with(JSONL_EXTENSION) {
                        #[cfg(feature = "debug")]
                        println!("Reading exported file: {}", path.display());
                        let path_clone = path.clone();

                        match read_jsonl_file(&path) {
                            Ok(records) => {
                                for record in records {
                                    // Use a HashMap to deduplicate records by name and cfg
                                    let key = if let Some(cfg) = &record.cfg {
                                        format!("{}#{}", record.name, cfg)
                                    } else {
                                        record.name.clone()
                                    };
                                    record_map.insert(key, record);
                                }
                            }
                            Err(e) => {
                                panic!("Failed to read {}: {}", path_clone.display(), e);
                            }
                        }
                    }
                }
            }
        }

        // Return deduplicated records for this group
        record_map.into_values().collect::<Vec<_>>()
    }

    /// Internal method to discover all available groups from the directory
    fn discover_groups<P: AsRef<Path>>(input_dir: P) -> HashSet<String> {
        let mut groups = HashSet::new();

        // Discover all available groups
        if let Ok(entries) = fs::read_dir(input_dir) {
            for entry in entries.flatten() {
                let path = entry.path();
                if let Some(file_name) = path.file_name().and_then(|n| n.to_str()) {
                    if file_name.ends_with(JSONL_EXTENSION) {
                        // Extract group name from filename (everything before the first underscore)
                        if let Some(underscore_pos) = file_name.find('_') {
                            let group_name = &file_name[..underscore_pos];
                            groups.insert(group_name.to_string());
                        }
                    }
                }
            }
        }

        groups
    }
}

/// Read the crate name from the stored file
fn read_stored_crate_name(input_dir: &Path) -> Option<String> {
    let crate_name_path = input_dir.join(CRATE_NAME_FILE);
    fs::read_to_string(crate_name_path)
        .ok()
        .map(|s| s.trim().to_string())
}

/// Read enabled features list from FEATURES_FILE and normalize into a sorted, deduplicated Vec<String>
fn read_features_from_out_dir(input_dir: &Path) -> Vec<String> {
    let features_path = input_dir.join(FEATURES_FILE);
    let Some(contents) = fs::read_to_string(features_path).ok() else {
        return Vec::new();
    };
    let mut features: Vec<String> = contents
        .lines()
        .map(str::trim)
        .filter(|l| !l.is_empty())
        .map(|s| s.to_string())
        .collect();
    features.sort();
    features.dedup();
    features
}

/// Builder for constructing a `Source` with custom options
pub struct Builder {
    input_dir: PathBuf,
    features_constant: Option<String>,
    target_triple: Option<String>,
    crate_name: Option<String>,
}

impl Builder {
    fn new<P: AsRef<Path>>(input_dir: P) -> Self {
        let target = std::env::var("TARGET").unwrap_or_default();
        let host = std::env::var("HOST").unwrap_or_default();
        let target_triple = if !target.is_empty() {
            Some(target)
        } else if !host.is_empty() {
            Some(host)
        } else {
            None
        };
        Self {
            input_dir: input_dir.as_ref().to_path_buf(),
            features_constant: Some("FEATURES".to_string()),
            target_triple,
            crate_name: None,
        }
    }

    /// Override the crate name recorded by the source crate's build script
    /// (`CARGO_PKG_NAME` at capture time) with the name **this** crate
    /// references the source crate by. Needed when the dependency is
    /// renamed in Cargo.toml:
    ///
    /// ```toml
    /// [dependencies]
    /// myflat = { package = "some-flat-crate", version = "..." }
    /// ```
    ///
    /// The name is stamped into every yielded item's `SourceLocation`, so
    /// it is what generated code qualifies references with
    /// (`myflat::…`) and what qualifies the features-guard constant. Being
    /// per-`Source`, it composes with multi-source registries — each
    /// chained stream carries its own (possibly overridden) origin.
    pub fn crate_name(mut self, name: impl Into<String>) -> Self {
        self.crate_name = Some(name.into());
        self
    }

    /// Enables or disables filtering by features when extracting collected data.
    /// Accepts name of the constant with the list of features in the source crate.
    ///
    /// Pass `None` to disable feature filtering.
    ///
    /// Pass, for example, `Some("source_crate::FEATURES")` or `Some("FEATURES")`
    /// to enable filtering. The value is the name of the features constant
    /// from the source crate.
    ///
    /// It's important to note that the set of features to filter is determined
    /// not by this constant, but by file "features.txt" in
    /// prebindgen output directory. These are features which the source crate was
    /// built with *as build.rs dependency*. The constant contains the features which
    /// the source crate was built *as library dependency*.
    /// The purpose of the constant is to use it in the assert in the
    /// generated code to ensure that both feature lists are the same.
    /// If it's not the case, compilation fails and it's developer's
    /// job to ensure that the source crate is configured in the same way for
    /// both `[dependencies]` and `[build-dependencies]`.
    ///
    /// Filtering is enabled by default; the default constant name is `FEATURES`.
    ///
    /// The features constant should be defined in the source crate as follows:
    /// ```rust
    /// const FEATURES: &str = prebindgen_proc_macro::features!();
    /// ```
    #[roxygen]
    pub fn enable_feature_filtering(
        mut self,
        /// Full name of the constant with features in the source crate
        name: Option<impl Into<String>>,
    ) -> Self {
        self.features_constant = name.map(|n| n.into());
        self
    }

    /// Enables filtering source code by target-triple. This is
    /// useful if some binding generators or custom tools over them
    /// encounter problems with build conditions in the source code.
    ///
    /// Accepts the target triple string.
    ///
    /// Default value is the value of the environment variable
    /// `TARGET` or `HOST` if `TARGET` is not set.
    ///
    /// Pass `None` to disable target filtering.
    ///
    /// When enabled, the code will be filtered by the target triple specified
    /// (architecture, vendor, os, env).
    ///
    /// It's useful to mention that it's hard to imagine a scenario
    /// when it's necessary to change the target triple.
    /// The `Source` object is used in the final crate, the one which builds the
    /// no-mangle library. On this stage the `TARGET` environment variable contains
    /// real target triple for cross-compilation scenarios.
    ///
    /// There is the case when build.rs is not aware of the cross-compilation
    /// target: the build.rs for library in [build-dependencies] section.
    /// This is the situation which the "ffi" crate found itself in. So sometimes
    /// for correct code generation it's necessary to inform the underlying "ffi"
    /// crate about the real target in a way like this:
    ///
    /// ```sh
    /// CROSS_TARGET=x86_64-pc-windows-gnu cargo build --target x86_64-pc-windows-gnu
    /// ```
    ///
    /// But in typical situation this doesn't affect the `Source` object which
    /// reads the data prepared by "ffi" crate on the following stage.
    ///
    /// So this `Some(<target-triple>)` is added mostly for API completeness,
    /// typical usage if this function will be with `None` argument to disable
    /// filtering by target triple if it's necessary to see the original
    /// collected code.
    #[roxygen]
    pub fn enable_target_filtering(
        mut self,
        /// The target triple to filter collected sources
        name: Option<impl Into<String>>,
    ) -> Self {
        self.target_triple = name.map(|n| n.into());
        self
    }

    /// Build the `Source` instance
    pub fn build(self) -> Source {
        Source::build_internal(
            &self.input_dir,
            self.features_constant,
            self.target_triple,
            self.crate_name,
        )
    }
}