ixa/

report.rs

use std::any::TypeId;
use std::cell::{RefCell, RefMut};
use std::env;
use std::fs::File;
use std::path::PathBuf;

use csv::Writer;
use serde::Serializer;

use crate::context::Context;
use crate::error::IxaError;
use crate::{define_data_plugin, error, trace, ContextBase, HashMap, HashMapExt};

// * file_prefix: precedes the report name in the filename. An example of a
// potential prefix might be scenario or simulation name
// * directory: location that the CSVs are written to. An example of this might
// be /data/
// * overwrite: if true, will overwrite existing files in the same location
pub struct ConfigReportOptions {
    pub file_prefix: String,
    pub output_dir: PathBuf,
    pub overwrite: bool,
}

impl ConfigReportOptions {
    #[must_use]
    pub fn new() -> Self {
        trace!("new ConfigReportOptions");
        // Sets the defaults
        ConfigReportOptions {
            file_prefix: String::new(),
            output_dir: env::current_dir().unwrap(),
            overwrite: false,
        }
    }
    /// Sets the file prefix option (e.g., "report_")
    pub fn file_prefix(&mut self, file_prefix: impl Into<String>) -> &mut ConfigReportOptions {
        let file_prefix = file_prefix.into();
        trace!("setting report prefix to {file_prefix}");
        self.file_prefix = file_prefix;
        self
    }
    /// Sets the directory where reports will be output
    pub fn directory(&mut self, directory: impl Into<PathBuf>) -> &mut ConfigReportOptions {
        let directory = directory.into();
        trace!("setting report directory to {directory:?}");
        self.output_dir = directory;
        self
    }
    /// Sets whether to overwrite existing reports of the same name if they exist
    pub fn overwrite(&mut self, overwrite: bool) -> &mut ConfigReportOptions {
        trace!("setting report overwrite {overwrite}");
        self.overwrite = overwrite;
        self
    }
}

impl Default for ConfigReportOptions {
    fn default() -> Self {
        Self::new()
    }
}

pub trait Report: 'static {
    // Returns report type
    fn type_id(&self) -> TypeId;
    // Serializes the data with the correct writer
    fn serialize(&self, writer: &mut Writer<File>);
}

/// # Errors
/// function will return Error if it fails to `serialize_str`
#[allow(dead_code)]
pub fn serialize_f64<S, const N: usize>(value: &f64, serializer: S) -> Result<S::Ok, S::Error>
where
    S: Serializer,
{
    let formatted = format!("{value:.N$}");
    serializer.serialize_str(&formatted)
}

/// # Errors
/// function will return Error if it fails to `serialize_str`
#[allow(dead_code)]
pub fn serialize_f32<S, const N: usize>(value: &f32, serializer: S) -> Result<S::Ok, S::Error>
where
    S: Serializer,
{
    let formatted = format!("{value:.N$}");
    serializer.serialize_str(&formatted)
}

struct ReportData {
    file_writers: RefCell<HashMap<TypeId, Writer<File>>>,
    config: ConfigReportOptions,
}

// Registers a data container that stores
// * file_writers: Maps report type to file writer
// * config: Contains all the customizable filename options that the user supplies
define_data_plugin!(
    ReportPlugin,
    ReportData,
    ReportData {
        file_writers: RefCell::new(HashMap::new()),
        config: ConfigReportOptions::new(),
    }
);

pub trait ContextReportExt: ContextBase {
    // Builds the filename. Called by `add_report`, `short_name` refers to the
    // report type. The three main components are `prefix`, `directory`, and
    // `short_name`.
    fn generate_filename(&mut self, short_name: &str) -> PathBuf {
        let data_container = self.get_data_mut(ReportPlugin);
        let prefix = &data_container.config.file_prefix;
        let directory = &data_container.config.output_dir;
        let short_name = short_name.to_string();
        let basename = format!("{prefix}{short_name}");
        directory.join(basename).with_extension("csv")
    }

    /// Add a report file keyed by a [`TypeId`].
    /// The `short_name` is used for file naming to distinguish what data each
    /// output file points to.
    /// # Errors
    /// If the file already exists and `overwrite` is set to false, raises an error and info message.
    /// If the file cannot be created, raises an error.
    fn add_report_by_type_id(&mut self, type_id: TypeId, short_name: &str) -> Result<(), IxaError> {
        trace!("adding report {short_name} by type_id {type_id:?}");
        let path = self.generate_filename(short_name);

        let data_container = self.get_data_mut(ReportPlugin);

        let file_creation_result = File::create_new(&path);
        let created_file = match file_creation_result {
            Ok(file) => file,
            Err(e) => match e.kind() {
                std::io::ErrorKind::AlreadyExists => {
                    if data_container.config.overwrite {
                        File::create(&path)?
                    } else {
                        error!("File already exists: {}. Please set `overwrite` to true in the file configuration and rerun.", path.display());
                        return Err(IxaError::IoError(e));
                    }
                }
                _ => {
                    return Err(IxaError::IoError(e));
                }
            },
        };
        let writer = Writer::from_writer(created_file);
        let mut file_writer = data_container.file_writers.borrow_mut();
        file_writer.insert(type_id, writer);
        Ok(())
    }

    /// Call `add_report` with each report type, passing the name of the report type.
    /// The `short_name` is used for file naming to distinguish what data each
    /// output file points to.
    /// # Errors
    /// If the file already exists and `overwrite` is set to false, raises an error and info message.
    /// If the file cannot be created, raises an error.
    fn add_report<T: Report + 'static>(&mut self, short_name: &str) -> Result<(), IxaError> {
        trace!("Adding report {short_name}");
        self.add_report_by_type_id(TypeId::of::<T>(), short_name)
    }

    /// Returns the writer registered for `type_id`.
    ///
    /// # Panics
    ///
    /// Panics if the report writers are already mutably borrowed or if no report
    /// was registered for `type_id`.
    fn get_writer(&self, type_id: TypeId) -> RefMut<Writer<File>> {
        // No data container will exist if no reports have been added
        let data_container = self.get_data(ReportPlugin);
        let writers = data_container
            .file_writers
            .try_borrow_mut()
            .expect("report writers are already mutably borrowed");
        RefMut::map(writers, |writers| {
            writers
                .get_mut(&type_id)
                .expect("no writer found for the report type; call add_report first")
        })
    }

    /// Writes a new row to the appropriate report file.
    ///
    /// # Panics
    ///
    /// Panics if the report writers are already mutably borrowed or if
    /// [`ContextReportExt::add_report`] was not called for `T`.
    fn send_report<T: Report>(&self, report: T) {
        let writer = &mut self.get_writer(report.type_id());
        report.serialize(writer);
    }

    /// Returns a `ConfigReportOptions` object which has setter methods for report configuration
    fn report_options(&mut self) -> &mut ConfigReportOptions {
        let data_container = self.get_data_mut(ReportPlugin);
        &mut data_container.config
    }
}
impl ContextReportExt for Context {}

#[cfg(test)]
mod test {
    use core::convert::TryInto;
    use std::thread;

    use serde_derive::{Deserialize, Serialize};
    use tempfile::tempdir;

    use super::*;
    use crate::{define_entity, define_property, define_report, info};

    define_entity!(Person);

    define_property!(
        struct IsRunner(bool),
        Person,
        default_const = IsRunner(false)
    );

    #[derive(Serialize, Deserialize)]
    struct SampleReport {
        id: u32,
        value: String,
    }

    define_report!(SampleReport);

    #[test]
    fn add_and_send_report() {
        let temp_dir = tempdir().unwrap();
        let path = PathBuf::from(&temp_dir.path());
        // We need the writer to go out of scope so the file is flushed
        {
            let mut context = Context::new();
            let config = context.report_options();
            config
                .file_prefix("prefix1_".to_string())
                .directory(path.clone());
            context.add_report::<SampleReport>("sample_report").unwrap();
            let report = SampleReport {
                id: 1,
                value: "Test Value".to_string(),
            };

            context.send_report(report);
        }

        let file_path = path.join("prefix1_sample_report.csv");
        assert!(file_path.exists(), "CSV file should exist");
        assert!(file_path.metadata().unwrap().len() > 0);

        let mut reader = csv::Reader::from_path(file_path).unwrap();
        for result in reader.deserialize() {
            let record: SampleReport = result.unwrap();
            assert_eq!(record.id, 1);
            assert_eq!(record.value, "Test Value");
        }
    }

    #[test]
    fn add_report_empty_prefix() {
        let temp_dir = tempdir().unwrap();
        let path = PathBuf::from(&temp_dir.path());
        // We need the writer to go out of scope so the file is flushed
        {
            let mut context = Context::new();
            let config = context.report_options();
            config.directory(path.clone());
            context.add_report::<SampleReport>("sample_report").unwrap();
            let report = SampleReport {
                id: 1,
                value: "Test Value".to_string(),
            };

            context.send_report(report);
        }
        let file_path = path.join("sample_report.csv");
        assert!(file_path.exists(), "CSV file should exist");
        assert!(file_path.metadata().unwrap().len() > 0);

        let mut reader = csv::Reader::from_path(file_path).unwrap();
        for result in reader.deserialize() {
            let record: SampleReport = result.unwrap();
            assert_eq!(record.id, 1);
            assert_eq!(record.value, "Test Value");
        }
    }

    struct PathBufWithDrop {
        file: PathBuf,
    }

    impl Drop for PathBufWithDrop {
        fn drop(&mut self) {
            std::fs::remove_file(&self.file).unwrap();
        }
    }

    #[test]
    fn add_report_no_dir() {
        // We need the writer to go out of scope so the file is flushed
        {
            let mut context = Context::new();
            let config = context.report_options();
            config.file_prefix("test_prefix_".to_string());
            context.add_report::<SampleReport>("sample_report").unwrap();
            let report = SampleReport {
                id: 1,
                value: "Test Value".to_string(),
            };

            context.send_report(report);
        }

        let path = env::current_dir().unwrap();
        let file_path = PathBufWithDrop {
            file: path.join("test_prefix_sample_report.csv"),
        };
        assert!(file_path.file.exists(), "CSV file should exist");
        assert!(file_path.file.metadata().unwrap().len() > 0);

        let mut reader = csv::Reader::from_path(&file_path.file).unwrap();
        for result in reader.deserialize() {
            let record: SampleReport = result.unwrap();
            assert_eq!(record.id, 1);
            assert_eq!(record.value, "Test Value");
        }
    }

    #[test]
    #[should_panic(expected = "no writer found for the report type; call add_report first")]
    fn send_report_without_adding_report() {
        let context = Context::new();
        let report = SampleReport {
            id: 1,
            value: "Test Value".to_string(),
        };

        context.send_report(report);
    }

    #[test]
    fn multiple_reports_one_context() {
        let temp_dir = tempdir().unwrap();
        let path = PathBuf::from(&temp_dir.path());
        // We need the writer to go out of scope so the file is flushed
        {
            let mut context = Context::new();
            let config = context.report_options();
            config
                .file_prefix("mult_report_".to_string())
                .directory(path.clone());
            context.add_report::<SampleReport>("sample_report").unwrap();
            let report1 = SampleReport {
                id: 1,
                value: "Value,1".to_string(),
            };
            let report2 = SampleReport {
                id: 2,
                value: "Value\n2".to_string(),
            };

            context.send_report(report1);
            context.send_report(report2);
        }

        let file_path = path.join("mult_report_sample_report.csv");
        assert!(file_path.exists(), "CSV file should exist");

        let mut reader = csv::Reader::from_path(file_path).expect("Failed to open CSV file");
        let mut records = reader.deserialize::<SampleReport>();

        let item1: SampleReport = records
            .next()
            .expect("No record found")
            .expect("Failed to deserialize record");
        assert_eq!(item1.id, 1);
        assert_eq!(item1.value, "Value,1");

        let item2: SampleReport = records
            .next()
            .expect("No second record found")
            .expect("Failed to deserialize record");
        assert_eq!(item2.id, 2);
        assert_eq!(item2.value, "Value\n2");
    }

    #[test]
    fn multithreaded_report_generation_thread_local() {
        let num_threads = 10;
        let num_reports_per_thread = 5;

        let mut handles = vec![];
        let temp_dir = tempdir().unwrap();
        let base_path = temp_dir.path().to_path_buf();

        for i in 0..num_threads {
            let path = base_path.clone();
            let handle = thread::spawn(move || {
                let mut context = Context::new();
                let config = context.report_options();
                config.file_prefix(i.to_string()).directory(path);
                context.add_report::<SampleReport>("sample_report").unwrap();

                for j in 0..num_reports_per_thread {
                    let report = SampleReport {
                        id: u32::try_from(i * num_reports_per_thread + j).unwrap(),
                        value: format!("Thread {i} Report {j}"),
                    };
                    context.send_report(report);
                }
            });

            handles.push(handle);
        }

        for handle in handles {
            handle.join().expect("Thread failed");
        }

        for i in 0..num_threads {
            let file_name = format!("{i}sample_report.csv");
            let file_path = base_path.join(file_name);
            assert!(file_path.exists(), "CSV file should exist");

            let mut reader = csv::Reader::from_path(file_path).expect("Failed to open CSV file");
            let records = reader.deserialize::<SampleReport>();

            for (j, record) in records.enumerate() {
                let record: SampleReport = record.expect("Failed to deserialize record");
                let id_expected = TryInto::<u32>::try_into(i * num_reports_per_thread + j).unwrap();
                assert_eq!(record.id, id_expected);
            }
        }
    }

    #[test]
    fn dont_overwrite_report() {
        let mut context1 = Context::new();
        let temp_dir = tempdir().unwrap();
        let path = PathBuf::from(&temp_dir.path());
        let config = context1.report_options();
        config
            .file_prefix("prefix1_".to_string())
            .directory(path.clone());
        context1
            .add_report::<SampleReport>("sample_report")
            .unwrap();
        let report = SampleReport {
            id: 1,
            value: "Test Value".to_string(),
        };

        context1.send_report(report);

        let file_path = path.join("prefix1_sample_report.csv");
        assert!(file_path.exists(), "CSV file should exist");

        let mut context2 = Context::new();
        let config = context2.report_options();
        config.file_prefix("prefix1_".to_string()).directory(path);
        info!("The next 'file already exists' error is intended for a passing test.");
        let result = context2.add_report::<SampleReport>("sample_report");
        assert!(result.is_err());
        let error = result.err().unwrap();
        match error {
            IxaError::IoError(e) => {
                assert_eq!(e.kind(), std::io::ErrorKind::AlreadyExists);
            }
            _ => {
                panic!("Unexpected error type");
            }
        }
    }

    #[test]
    fn overwrite_report() {
        let mut context1 = Context::new();
        let temp_dir = tempdir().unwrap();
        let path = PathBuf::from(&temp_dir.path());
        let config = context1.report_options();
        config
            .file_prefix("prefix1_".to_string())
            .directory(path.clone());
        context1
            .add_report::<SampleReport>("sample_report")
            .unwrap();
        let report = SampleReport {
            id: 1,
            value: "Test Value".to_string(),
        };

        context1.send_report(report);

        let file_path = path.join("prefix1_sample_report.csv");
        assert!(file_path.exists(), "CSV file should exist");

        let mut context2 = Context::new();
        let config = context2.report_options();
        config
            .file_prefix("prefix1_".to_string())
            .directory(path)
            .overwrite(true);
        let result = context2.add_report::<SampleReport>("sample_report");
        assert!(result.is_ok());
        let file = File::open(file_path).unwrap();
        let reader = csv::Reader::from_reader(file);
        let records = reader.into_records();
        assert_eq!(records.count(), 0);
    }
}