scyros 0.3.1

// Copyright 2025 Andrea Gilot
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Utility functions for file operations and I/O.

use anyhow::{bail, Context, Error, Result};
use pathdiff::diff_paths;
use polars::io::SerWriter;
use polars::prelude::{CsvReadOptions, CsvWriter, Schema};
use polars::{frame::DataFrame, io::SerReader};
use walkdir::WalkDir;

use std::fs;
use std::io::BufWriter;
use std::path::{Component, PathBuf};
use std::sync::Arc;
use std::{
    fs::File,
    io::{BufRead, BufReader, Lines},
    path::Path,
};

#[derive(PartialEq, Eq, Debug, Copy, Clone)]
pub enum FileMode {
    Read,
    Overwrite,
    Append,
}

/// Opens a file. In overwrite or append mode, creates the file if it does not exist.
///
/// # Arguments
///
/// * `path` - The path to the file.
/// * `mode` - The mode to open the file in.
///
/// # Returns
///
/// A file in the specified mode or an error if the file could not be opened or created.
pub fn open_file(path: impl AsRef<Path>, mode: FileMode) -> Result<File> {
    if let Some(parent) = path.as_ref().parent() {
        if let Some(parent_path) = parent.to_str() {
            create_dir(parent_path)?;
        }
    }
    match mode {
        FileMode::Read => std::fs::File::open(&path),
        FileMode::Overwrite => std::fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(&path),
        FileMode::Append => std::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(&path),
    }
    .with_context(|| format!("Could not open {}", &path.as_ref().display()))
}

pub fn check_path(path: &str) -> Result<PathBuf> {
    if Path::new(path).exists() {
        Ok(PathBuf::from(path))
    } else {
        bail!("File or directory {path} not found")
    }
}

// TODO : Test
/// Loads a file into memory if its size is less than a given limit.
///
/// # Arguments
///
/// * `path` - The path to the file.
/// * `memory_limit` - The maximum size of the file in bytes.
///
/// # Returns
///
/// An array of bytes containing the content of the file or an error if the file could not be read.
/// Two kinds of errors are possible:
/// * If the file size exceeds the memory limit, returns the size of the file.
/// * If the file could not be read, returns an error.
pub fn load_file(
    path: impl AsRef<Path>,
    memory_limit: u64,
) -> Result<core::result::Result<Vec<u8>, u64>> {
    let metadata = std::fs::metadata(&path).with_context(|| {
        format!(
            "Could not fetch metadata for file {}",
            &path.as_ref().display()
        )
    })?;
    let file_size = metadata.len();
    if file_size > memory_limit {
        Ok(Err(file_size))
    } else {
        std::fs::read(&path)
            .map(Ok)
            .with_context(|| format!("Could not read file {}", &path.as_ref().display()))
    }
}

/// Returns an iterator on the lines of a file.
///
/// # Arguments
///
/// * `path` - The path to the file to read the lines from.
///
/// # Returns
///
/// A vector containing all the lines of the file or an error if the file could not be read.
pub fn file_lines(path: impl AsRef<Path>) -> Result<Lines<BufReader<File>>, Error> {
    Ok(std::io::BufReader::new(open_file(path, FileMode::Read)?).lines())
}

/// Counts the number of lines in a file.
///
/// # Arguments
///
/// * `path` - The path to the file.
pub fn file_lines_count(path: impl AsRef<Path>) -> Result<usize, Error> {
    Ok(file_lines(path)?.count())
}

/// Creates a directory without returning an error if it already exists or one of its parents does not exist.
/// If one of the parents does not exist, creates it as well.
///
/// # Arguments
///
/// * `path` - The path to the directory.
///
/// # Returns
///
/// An error if the directory could not be created.
pub fn create_dir(path: impl AsRef<Path>) -> Result<(), Error> {
    let path_buf = path.as_ref().to_path_buf();
    match std::fs::create_dir_all(&path_buf) {
        Ok(_) => Ok(()),
        Err(e) => {
            if e.kind() != std::io::ErrorKind::AlreadyExists {
                bail!(format!(
                    "Could not create directory {}: {}",
                    path_buf.display(),
                    e
                ))
            } else {
                Ok(())
            }
        }
    }
}

/// Deletes a directory.
///
/// # Arguments
///
/// * `path` - The path to the directory.
/// * `silent` - If true, does not raise an error if the directory does not exist.
///
/// # Returns
///
/// An error if the directory could not be deleted.
pub fn delete_dir(path: impl AsRef<Path>, silent: bool) -> Result<()> {
    let path_buf = path.as_ref().to_path_buf();
    match std::fs::remove_dir_all(&path_buf) {
        Ok(_) => Ok(()),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound && silent => Ok(()),
        Err(e) => Err(e.into()),
    }
}

/// Deletes a file.
///
/// # Arguments
///
/// * `path` - The path to the file.
/// * `silent` - If true, does not raise an error if the file does not exist.
///
/// # Returns
///
/// An error if the file could not be deleted.
///
pub fn delete_file(path: impl AsRef<Path>, silent: bool) -> Result<()> {
    let path_buf = path.as_ref().to_path_buf();
    match std::fs::remove_file(&path_buf) {
        Ok(_) => Ok(()),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound && silent => Ok(()),
        Err(e) => bail!(format!(
            "Could not delete file {}: {}",
            path_buf.display(),
            e
        )),
    }
}

/// Writes content to a file, creating the file and its parent directories if they do not exist.
///
/// # Arguments
/// * `path` - The path to the file.
/// * `content` - The content to write to the file.
///
/// # Returns
/// An error if the file could not be written.
pub fn write_file(path: impl AsRef<Path>, content: impl AsRef<[u8]>) -> Result<()> {
    if let Some(parent) = path.as_ref().parent() {
        create_dir(parent)?;
    }
    fs::write(&path, content)?;

    Ok(())
}

/// Reads a CSV file into a DataFrame.
///
/// # Arguments
/// * `path` - The path to the CSV file.
/// * `schema` - A schema, i.e., a list of column names and their associated data types.
///   This list does not need to contain all the columns of the CSV file, nor does it need to strictly contain columns of the CSV file.
///   The columns of the CSV file that are not in the schema will be read with an inferred data type.
/// * `columns` - A list of column names to read from the CSV file. If None, reads all columns.
///
/// # Returns
/// A DataFrame containing the data from the CSV file or an error if the file could not be read or if the data could not be parsed according to the schema.
pub fn open_csv(
    path: &str,
    schema: Option<Schema>,
    columns: Option<Vec<&str>>,
) -> Result<DataFrame, Error> {
    CsvReadOptions::default()
        .with_columns(
            columns.map(|cols| Arc::from(cols.into_iter().map(|s| s.into()).collect::<Vec<_>>())),
        )
        .with_schema_overwrite(schema.map(Arc::new))
        .with_has_header(true)
        .into_reader_with_file_handle(BufReader::new(open_file(path, FileMode::Read)?))
        .finish()
        .with_context(|| format!("Could not read {path}"))
}

/// Writes a DataFrame to a CSV file.
///
/// # Arguments
/// * `path` - The path to the output CSV file.
/// * `df` - The DataFrame to write to the CSV file.
///
/// # Returns
/// An error if the DataFrame could not be written to the CSV file.
pub fn write_csv(path: &str, df: &mut DataFrame) -> Result<()> {
    CsvWriter::new(BufWriter::new(open_file(path, FileMode::Overwrite)?))
        .include_header(true)
        .with_separator(b',')
        .finish(df)
        .with_context(|| format!("Could not write to {path}"))
}

pub fn is_empty_dir(path: impl AsRef<Path>) -> Result<bool> {
    Ok(fs::read_dir(path)?.next().is_none())
}

pub fn delete_empty_dirs(path: impl AsRef<Path>) -> Result<()> {
    for entry in WalkDir::new(path)
        .contents_first(true)
        .into_iter()
        .filter_map(Result::ok)
        .filter(|e| e.file_type().is_dir())
    {
        if is_empty_dir(entry.path())? {
            fs::remove_dir(entry.path())?;
        }
    }
    Ok(())
}

/// Returns a list of files with a given extension in a directory and its subdirectories,
/// sorted by their proximity to a pivot file.
/// The proximity is defined as the number of directory levels to go up from the pivot file
/// to reach the common ancestor directory, and then the number of directory levels to go down
/// to reach the file.
///
/// # Arguments
///
/// * `root_dir` - The root directory to search for files.
/// * `pivot_file` - The pivot file to measure the proximity from.
/// * `ext` - The file extension to filter the files by (case insensitive).
///
/// # Returns
///
/// A vector of paths to the files with the given extension, sorted by their proximity to the pivot file,
/// or an error if the root directory or the pivot file do not exist, or if the pivot file is not in the root directory.
///
/// # Examples
/// If the directory structure is as follows:
/// ```text
/// src/
/// ├── main.rs
/// ├── utils/
/// │   ├── snippets/
/// │   │   ├── example.rs
/// │   │   └── example.py
/// │   ├── foo.rs
/// │   ├── bar.rs
/// │   └── bar.py
/// └── io/
///     └── fs.rs
/// ```
/// and the pivot file is `src/utils/foo.rs`, calling
/// `files_sorted_by_proximity("src", "src/utils/foo.rs", "rs")`
/// will return the files in the following order:
/// 1. `src/utils/foo.rs` (0 ups, 0 downs)
/// 2. `src/utils/bar.rs` (1 up, 1 down)
/// 3. `src/utils/snippets/example.rs` (1 up, 2 downs)
/// 4. `src/main.rs` (2 ups, 1 down)
/// 5. `src/io/fs.rs` (2 ups, 2 down)
pub fn files_sorted_by_proximity(
    root_dir: impl AsRef<Path>,
    pivot_file: impl AsRef<Path>,
    ext: &str,
) -> Result<Vec<PathBuf>, Error> {
    let pivot_file = pivot_file.as_ref();
    let root_dir = root_dir.as_ref();

    if !pivot_file.exists() {
        bail!("Pivot file {pivot_file:?} does not exist")
    } else {
        let pivot_canon = pivot_file
            .canonicalize()
            .with_context(|| format!("Could not canonicalize pivot file {pivot_file:?}"))?;
        let root_canon = root_dir
            .canonicalize()
            .with_context(|| format!("Could not canonicalize root dir {root_dir:?}"))?;

        if !pivot_canon.starts_with(&root_canon) {
            bail!("Pivot file {pivot_file:?} is not in root dir {root_dir:?}")
        } else {
            let mut files: Vec<PathBuf> = WalkDir::new(root_dir)
                .into_iter()
                .filter_map(Result::ok)
                .filter(|e| e.file_type().is_file())
                .map(|e| e.into_path())
                .filter(|p| {
                    p.extension()
                        .and_then(|e| e.to_str())
                        .map(|e| e.eq_ignore_ascii_case(ext))
                        .unwrap_or(false)
                })
                .collect();

            files.sort_by_key(|p| {
                // Safe unwrap because pivot_file is guaranteed to be a file
                let rel: PathBuf = diff_paths(p, pivot_file).unwrap();

                let mut ups = 0;
                let mut total = 0;
                for comp in rel.components() {
                    if matches!(comp, Component::ParentDir) {
                        ups += 1;
                    } else if !matches!(comp, Component::CurDir) {
                        total += 1;
                    }
                }
                (ups, total)
            });

            Ok(files)
        }
    }
}

#[cfg(test)]
mod io_tests {

    use std::io::Write;
    use std::path::Path;

    use anyhow::{ensure, Ok};

    use super::*;

    #[test]
    fn read_file_test() -> Result<()> {
        let file = open_file("tests/data/non_existent_file.txt", FileMode::Read);
        ensure!(file.is_err());

        open_file("tests/data/empty.csv", FileMode::Read)?;
        Ok(())
    }

    #[test]
    fn file_lines_test() -> Result<()> {
        let path = "tests/data/small_file.csv";
        let mut lines = file_lines(path)?;
        assert_eq!(lines.next().unwrap()?, "id,name,fork");
        assert_eq!(lines.next().unwrap()?, "0,a,1");
        assert_eq!(lines.next().unwrap()?, "1,b,0");
        assert_eq!(lines.next().unwrap()?, "2,c,1");
        assert_eq!(lines.next().unwrap()?, "3,d,0");
        ensure!(lines.next().is_none());
        Ok(())
    }

    #[test]
    fn create_delete_dir_test() -> Result<()> {
        let test_dir = "tests";
        create_dir(test_dir)?;

        delete_dir(format!("{test_dir}/new_dir"), true)?;
        ensure!(delete_dir(format!("{test_dir}/new_dir"), false).is_err());

        let new_dir = format!("{test_dir}/new_dir/new_dir");
        ensure!(!Path::new(&new_dir).exists());
        create_dir(&new_dir)?;
        ensure!(Path::new(&new_dir).exists());

        delete_dir(format!("{test_dir}/new_dir"), false)?;
        ensure!(!Path::new(&new_dir).exists());
        Ok(())
    }

    #[test]
    fn create_delete_file_test() -> Result<()> {
        let test_file = "tests/new_file.txt";

        delete_file(test_file, true)?;
        ensure!(delete_file(test_file, false).is_err());

        ensure!(!Path::new(&test_file).exists());

        open_file(test_file, FileMode::Overwrite)?;

        ensure!(Path::new(&test_file).exists());

        delete_file(test_file, false)?;

        ensure!(!Path::new(&test_file).exists());
        Ok(())
    }

    #[test]
    fn write_file_test() -> Result<()> {
        let path = "tests/data/abc.txt";

        {
            let file = open_file(path, FileMode::Overwrite)?;
            write!(&file, "abc")?;
        }

        let content = std::fs::read_to_string(path)?;
        let lines: Vec<&str> = content.lines().collect();
        ensure!(lines.len() == 1);
        assert_eq!(lines[0], "abc");

        {
            let file = open_file(path, FileMode::Append)?;
            write!(&file, "okok")?;
        }

        let content = std::fs::read_to_string(path)?;
        let lines: Vec<&str> = content.lines().collect();
        ensure!(lines.len() == 1);
        assert_eq!(lines[0], "abcokok");

        {
            let file = open_file(path, FileMode::Overwrite)?;
            write!(&file, "abc")?;
        }

        let content = std::fs::read_to_string(path)?;
        let lines: Vec<&str> = content.lines().collect();
        ensure!(lines.len() == 1);
        assert_eq!(lines[0], "abc");
        Ok(())
    }

    #[test]
    fn line_count_test() -> Result<()> {
        let count = file_lines_count("tests/data/small_file.csv")?;
        assert_eq!(count, 5);

        ensure!(file_lines_count("tests/data/non_existent_file.csv").is_err());
        Ok(())
    }

    #[test]
    fn files_sorted_by_proximity_test() -> Result<()> {
        let root_dir = "tests/data/test_project";
        let pivot_file = "tests/data/test_project/utils/foo.rs";

        let files = files_sorted_by_proximity(root_dir, pivot_file, "rs")?;
        let files = files
            .into_iter()
            .map(|p| p.to_str().unwrap().to_string())
            .collect::<Vec<_>>();
        let expected_files = vec![
            "tests/data/test_project/utils/foo.rs",
            "tests/data/test_project/utils/bar.rs",
            "tests/data/test_project/utils/snippets/example.rs",
            "tests/data/test_project/main.rs",
            "tests/data/test_project/io/fs.rs",
        ];
        assert_eq!(files, expected_files);
        Ok(())
    }
}