mj_io 0.1.7

/*
Some helpful utilities for i/o operations. Local files only!


1. Listing files in a directory 
2. Reading files from s3/local into a vector of bytes
3. Writing a vector of bytes to s3/local
4. Compressing data


A note on extensions:
We will ONLY ever be concerned about files that have extensions of:
- .jsonl --> uncompressed, each line is a json string with a field 'text'
- .jsonl.gz -> jsonl compressed with gz
- .jsonl.zstd | .jsonl.zst -> jsonl compressed with zstandard 

Compression schemes will always be inferred from extension
*/

use std::io::BufWriter;
use std::io::BufRead;
use std::fs::create_dir_all;
use std::fs::File;
use anyhow::Error;
use anyhow::anyhow;
use std::path::PathBuf;
use glob::glob;
use flate2::read::MultiGzDecoder;
use zstd::stream::read::Decoder as ZstdDecoder;
use std::io::{BufReader, Cursor, Write, Read};
use flate2::write::GzEncoder;
use flate2::Compression;
use zstd::stream::write::Encoder as ZstdEncoder;
use indicatif::{ProgressBar, ProgressStyle};

const VALID_EXTS: &[&str] = &[".jsonl", ".jsonl.gz", ".jsonl.zstd", ".jsonl.zst", ".json.gz", ".json.zst"];

/*======================================================================
=                               Progress Bar stuff                     =
======================================================================*/



pub fn build_pbar(num_items: usize, units: &str) -> ProgressBar {
    let mut template = String::from(units);
    template.push_str(" {human_pos}/{human_len} [{elapsed_precise}/{duration_precise}] [{wide_bar:.cyan/blue}]");
    let pbar = ProgressBar::new(num_items as u64)
        .with_style(
            ProgressStyle::with_template(&template).unwrap()
        );
    pbar.inc(0);
    pbar
}




/*======================================================================
=                              Listing files                           =
======================================================================*/




pub fn expand_dirs(paths: Vec<PathBuf>, manual_ext: Option<&[&str]>) -> Result<Vec<PathBuf>, Error> {
    // For local directories -> does a glob over each directory to get all files with given extension
    let exts = if !manual_ext.is_none() {
    	manual_ext.unwrap()
    } else {
    	VALID_EXTS
    };
    let mut files: Vec<PathBuf> = Vec::new();
    for path in paths {
        if path.is_dir() {
            let path_str = path
                .to_str()
                .ok_or_else(|| anyhow!("invalid path '{}'", path.to_string_lossy()))?;
        	for ext in exts {
        		let pattern = format!("{}/**/*{}", path_str, ext);
        		for entry in glob(&pattern).expect("Failed to read glob pattern") {
        			if let Ok(path) = entry {
        				files.push(path)
        			}
        		}
        	}
        } else {
            files.push(path.clone());
        }
    }
    Ok(files)
}


pub fn has_json_extension(path: &PathBuf) -> bool {
    if let Some(ext) = path.extension() {
        if ext == "json" {
            return true;
        } else if let Some(ext_os_str) = path.extension() {
            if let Some(ext_str) = ext_os_str.to_str() {
                return ext_str.starts_with("json.");
            }
        }
    }
    false
}

/*====================================================================
=                            Naming files                            =
====================================================================*/

pub fn get_output_filename(input_path: &PathBuf, config_input_dir: &PathBuf, config_output_dir: &PathBuf) -> Result<PathBuf, Error> {
    // Given an input path that starts with config_input_dir, replaces the input_dir w/ the output_dir
    // NOTE: Will explode if input_path does not start with input_dir
    let replaced = input_path.clone()
        .strip_prefix(config_input_dir)
        .ok()
        .map(|stripped| config_output_dir.clone().join(stripped)).unwrap();
    Ok(replaced)
}




/*====================================================================
=                           Reading files                            =
====================================================================*/

// Enum to handle different reader types
pub enum FileReader {
    Memory(BufReader<Cursor<Vec<u8>>>),
    Stream(Box<dyn BufRead>),
}

impl FileReader {
    // Provide access to lines iterator
    pub fn lines(self) -> Box<dyn Iterator<Item = std::io::Result<String>>> {
        match self {
            FileReader::Memory(reader) => Box::new(reader.lines()),
            FileReader::Stream(reader) => Box::new(reader.lines()),
        }
    }
    
    // For cases where you need direct BufRead access
    pub fn as_buf_read(&mut self) -> &mut dyn BufRead {
        match self {
            FileReader::Memory(ref mut reader) => reader,
            FileReader::Stream(ref mut reader) => &mut **reader,
        }
    }
}

// Original function for backward compatibility - always loads into memory
pub fn read_pathbuf_to_mem(input_file: &PathBuf) -> Result<BufReader<Cursor<Vec<u8>>>, Error> {
    // Generic method to read local or s3 file into memory
    let contents = read_local_file_into_memory(input_file).expect("Failed to read contents into memory");
    let reader = BufReader::new(contents);
    Ok(reader)
}

// New function that can either stream or load into memory
pub fn read_pathbuf(input_file: &PathBuf, stream: bool) -> Result<FileReader, Error> {
    if stream {
        read_pathbuf_stream(input_file)
    } else {
        let mem_reader = read_pathbuf_to_mem(input_file)?;
        Ok(FileReader::Memory(mem_reader))
    }
}

// Streaming version - doesn't load entire file into memory
fn read_pathbuf_stream(input_file: &PathBuf) -> Result<FileReader, Error> {
    let file = File::open(input_file)?;
    let ext = input_file.extension()
        .and_then(|e| e.to_str())
        .map(|e| e.to_lowercase())
        .unwrap_or_default();
    
    let reader: Box<dyn BufRead> = match ext.as_str() {
        "gz" => {
            let decoder = MultiGzDecoder::new(file);
            Box::new(BufReader::new(decoder))
        },
        "zstd" | "zst" => {
            let decoder = ZstdDecoder::new(file)?;
            Box::new(BufReader::new(decoder))
        },
        _ => {
            Box::new(BufReader::new(file))
        }
    };
    
    Ok(FileReader::Stream(reader))
}

fn read_local_file_into_memory(input_file: &PathBuf) -> Result<Cursor<Vec<u8>>, Error> {
    let mut file = File::open(input_file)?;
    let mut contents = Vec::new();
    let ext = input_file.extension().unwrap().to_string_lossy().to_lowercase();
    
    if ext == "gz" {
        let mut decoder = MultiGzDecoder::new(file);
        decoder.read_to_end(&mut contents)?;
    } else if ext == "zstd" || ext == "zst" {
        let mut decoder = ZstdDecoder::new(file)?;
        // Updates: if incomplete frames on zst files, read as buffer first
        match decoder.read_to_end(&mut contents) { // <- Try the original thing
            Ok(_) => {}, // Success, we're done
            Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => { // <-- eof error leads to bufreads
                // Fallback: re-open and do streaming read for incomplete files
                let file = File::open(input_file)?;
                let mut decoder = ZstdDecoder::new(file)?;
                contents.clear();
                
                let mut buffer = [0; 64 * 1024]; 
                loop {
                    match decoder.read(&mut buffer) {
                        Ok(0) => break,
                        Ok(n) => contents.extend_from_slice(&buffer[..n]),
                        Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => break,
                        Err(e) => return Err(e.into()),
                    }
                }
            }
            Err(e) => return Err(e.into()),
        }
    } else {
        file.read_to_end(&mut contents)?;
    }
    
    Ok(Cursor::new(contents))
}

/*====================================================================
=                          Writing files                             =
====================================================================*/


// Writer enum to handle different writer types
pub enum FileWriter {
    Plain(BufWriter<File>),
    Gzip(Option<GzEncoder<File>>),
    Zstd(Option<ZstdEncoder<'static, File>>),
}

impl Drop for FileWriter {
    fn drop(&mut self) {
        // Best effort to finish the writer - ignore errors since we can't return them from drop
        match self {
            FileWriter::Plain(w) => {
                let _ = w.flush();
            },
            FileWriter::Gzip(ref mut w) => {
                if let Some(encoder) = w.take() {
                    let _ = encoder.finish();
                }
            },
            FileWriter::Zstd(ref mut w) => {
                if let Some(encoder) = w.take() {
                    let _ = encoder.finish();
                }
            },
        }
    }
}

impl FileWriter {
    // Write a line (adds newline if not present)
    pub fn write_line(&mut self, line: &Vec<u8>) -> Result<(), std::io::Error> {
        let final_byte = line[line.len() - 1];

        match self {
            FileWriter::Plain(w) => {
                w.write_all(line)?;
                if final_byte != b'\n' {
                    w.write_all(b"\n")?;
                }
                Ok(())
            },
            FileWriter::Gzip(ref mut w) => {
                if let Some(encoder) = w.as_mut() {
                    encoder.write_all(line)?;
                    if final_byte != b'\n' {
                        encoder.write_all(b"\n")?;
                    }
                    Ok(())
                } else {
                    Err(std::io::Error::new(std::io::ErrorKind::Other, "Writer already finished"))
                }
            },
            FileWriter::Zstd(ref mut w) => {
                if let Some(encoder) = w.as_mut() {
                    encoder.write_all(line)?;
                    if final_byte != b'\n' {
                        encoder.write_all(b"\n")?;
                    }
                    Ok(())
                } else {
                    Err(std::io::Error::new(std::io::ErrorKind::Other, "Writer already finished"))
                }
            },
        }
    }
    
    // Write raw bytes
    pub fn write_all(&mut self, bytes: &[u8]) -> Result<(), std::io::Error> {
        match self {
            FileWriter::Plain(w) => w.write_all(bytes),
            FileWriter::Gzip(ref mut w) => {
                if let Some(encoder) = w.as_mut() {
                    encoder.write_all(bytes)
                } else {
                    Err(std::io::Error::new(std::io::ErrorKind::Other, "Writer already finished"))
                }
            },
            FileWriter::Zstd(ref mut w) => {
                if let Some(encoder) = w.as_mut() {
                    encoder.write_all(bytes)
                } else {
                    Err(std::io::Error::new(std::io::ErrorKind::Other, "Writer already finished"))
                }
            },
        }
    }
    
    // Flush the buffer
    pub fn flush(&mut self) -> Result<(), std::io::Error> {
        match self {
            FileWriter::Plain(w) => w.flush(),
            FileWriter::Gzip(ref mut w) => {
                if let Some(encoder) = w.as_mut() {
                    encoder.flush()
                } else {
                    Err(std::io::Error::new(std::io::ErrorKind::Other, "Writer already finished"))
                }
            },
            FileWriter::Zstd(ref mut w) => {
                if let Some(encoder) = w.as_mut() {
                    encoder.flush()
                } else {
                    Err(std::io::Error::new(std::io::ErrorKind::Other, "Writer already finished"))
                }
            },
        }
    }
    
    // Finish writing and close the file (consumes the writer)
    pub fn finish(mut self) -> Result<(), std::io::Error> {
        match self {
            FileWriter::Plain(ref mut w) => w.flush(),
            FileWriter::Gzip(ref mut w) => {
                if let Some(encoder) = w.take() {
                    encoder.finish()?;
                    Ok(())
                } else {
                    Err(std::io::Error::new(std::io::ErrorKind::Other, "Writer already finished"))
                }
            },
            FileWriter::Zstd(ref mut w) => {
                if let Some(encoder) = w.take() {
                    encoder.finish()?;
                    Ok(())
                } else {
                    Err(std::io::Error::new(std::io::ErrorKind::Other, "Writer already finished"))
                }
            },
        }
    }
}

// Create a streaming writer for a file
pub fn create_writer(output_file: &PathBuf) -> Result<FileWriter, Error> {
    // Create parent directories if they don't exist
    if let Some(parent_dir) = output_file.parent() {
        if !parent_dir.exists() {
            create_dir_all(parent_dir)?;
        }
    }
    
    let file = File::create(output_file)?;
    let ext = output_file.extension()
        .and_then(|e| e.to_str())
        .map(|e| e.to_lowercase())
        .unwrap_or_default();
    
    let writer = match ext.as_str() {
        "gz" => {
            FileWriter::Gzip(Some(GzEncoder::new(file, Compression::default())))
        },
        "zstd" | "zst" => {
            FileWriter::Zstd(Some(ZstdEncoder::new(file, 0)?))
        },
        _ => {
            FileWriter::Plain(BufWriter::new(file))
        }
    };
    
    Ok(writer)
}

// Helper function to write an iterator of lines to a file
pub fn write_lines_to_file<I, S>(output_file: &PathBuf, lines: I) -> Result<(), Error> 
where
    I: Iterator<Item = S>,
    S: AsRef<Vec<u8>>,
{
    let mut writer = create_writer(output_file)?;
    for line in lines {
        writer.write_line(line.as_ref())?;
    }
    writer.finish()?;
    Ok(())
}


pub fn write_mem_to_pathbuf(contents: &[u8], output_file: &PathBuf) -> Result<(), Error> {
	let compressed_data = compress_data(contents.to_vec(), output_file);
    if let Some(parent_dir) = output_file.parent() {
            if !parent_dir.exists() {
                create_dir_all(parent_dir).unwrap()
             }
        let mut file = File::create(output_file).expect(format!("Unable to create output file {:?}", output_file).as_str());
        file.write_all(&compressed_data).expect(format!("Unable to write to {:?}", output_file).as_str());

    }
    Ok(())
}



fn compress_data(data: Vec<u8>, filename: &PathBuf) -> Vec<u8> {
    // Given a filename with an extension, compresses a bytestream accordingly 
    // {zst, zstd} -> zstandard, {gz} -> gzip, anything else -> nothing
    let output_data = match filename.extension().unwrap().to_str() {
        Some("gz") => {
            let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
            encoder.write_all(&data).unwrap();            
            encoder.finish().unwrap()
        },
        Some("zstd") | Some("zst") => {
            let mut encoder = ZstdEncoder::new(Vec::new(), 0).unwrap();
            encoder.write_all(&data).unwrap();            
            encoder.finish().unwrap()
        },
        _ => {data}
    };
    output_data
}