arcweight 0.3.0

//! FST Archive (FAR) format for storing multiple FSTs.
//!
//! This module provides support for FST Archives (FAR), a container format that
//! stores multiple named FSTs in a single file. This is analogous to tar archives
//! for files, enabling efficient storage and retrieval of FST collections.
//!
//! # Use Cases
//!
//! FAR files are commonly used for:
//!
//! - **Language models:** Store n-gram models for different contexts
//! - **Pronunciation lexicons:** Multiple dictionaries (formal, casual, domain-specific)
//! - **Rule sets:** Phonological or morphological rule collections
//! - **Model versions:** Keep multiple model versions in a single file
//! - **Parallel data:** Store aligned FSTs for source/target languages
//!
//! # Format Overview
//!
//! The FAR format consists of:
//!
//! ```text
//! +----------------------------------+
//! | FST 1 (OpenFST binary format)    |
//! +----------------------------------+
//! | FST 2 (OpenFST binary format)    |
//! +----------------------------------+
//! | ...                              |
//! +----------------------------------+
//! | Index:                           |
//! |   num_entries (4B)               |
//! |   entry 1: name_len, name, offset|
//! |   entry 2: name_len, name, offset|
//! |   ...                            |
//! +----------------------------------+
//! | Index position (8B)              |
//! +----------------------------------+
//! ```
//!
//! The index at the end enables efficient random access to FSTs by name
//! without scanning the entire file.
//!
//! # Examples
//!
//! ## Writing a FAR File
//!
//! ```no_run
//! use arcweight::prelude::*;
//! use arcweight::io::{create_far, FarWriter};
//!
//! // Create FAR writer
//! let mut writer = create_far("models.far")?;
//!
//! // Add multiple FSTs
//! let fst1 = VectorFst::<TropicalWeight>::new();
//! let fst2 = VectorFst::<TropicalWeight>::new();
//!
//! writer.add("model_v1", &fst1)?;
//! writer.add("model_v2", &fst2)?;
//!
//! // Finalize (writes index)
//! writer.finish()?;
//! # Ok::<(), Box<dyn std::error::Error>>(())
//! ```
//!
//! ## Reading from a FAR File
//!
//! ```no_run
//! use arcweight::prelude::*;
//! use arcweight::io::open_far;
//!
//! // Open FAR reader
//! let mut reader = open_far("models.far")?;
//!
//! // List all FSTs in the archive
//! for name in reader.list() {
//!     println!("Found FST: {}", name);
//! }
//!
//! // Read specific FST by name
//! if let Some(fst) = reader.read("model_v1")? {
//!     println!("Loaded FST with {} states", fst.num_states());
//! }
//! # Ok::<(), Box<dyn std::error::Error>>(())
//! ```
//!
//! # Performance
//!
//! | Operation | Complexity |
//! |-----------|------------|
//! | Open archive | O(n) where n = number of entries |
//! | List entries | O(1) |
//! | Read by name | O(1) seek + O(FST size) read |
//! | Add FST | O(FST size) |
//! | Finish | O(n) where n = number of entries |
//!
//! # Limitations
//!
//! - Currently only supports `TropicalWeight` FSTs (uses OpenFST format internally)
//! - FST names must be valid UTF-8 strings
//! - No compression (FSTs stored in raw OpenFST format)
//! - No support for streaming iteration (must read by name)
//!
//! # References
//!
//! - OpenFST FAR format: <https://www.openfst.org/twiki/bin/view/FST/FstExtensions>

use crate::fst::VectorFst;
use crate::semiring::TropicalWeight;
use crate::Result;
use byteorder::{LittleEndian, WriteBytesExt};
use std::collections::HashMap;
use std::io::{Read, Seek, Write};

/// Reader for FST Archive (FAR) files.
///
/// `FarReader` provides random access to FSTs stored in a FAR archive by name.
/// The index is read once when the reader is created, enabling O(1) lookups.
///
/// # Type Parameters
///
/// * `R` - The underlying reader type, must implement `Read + Seek`
///
/// # Examples
///
/// ```no_run
/// use arcweight::io::{open_far, FarReader};
/// use arcweight::Fst;
/// use std::fs::File;
/// use std::io::BufReader;
///
/// // Open from file path (convenience function)
/// let mut reader = open_far("archive.far")?;
///
/// // Or create from any Read + Seek type
/// let file = File::open("archive.far")?;
/// let mut reader = FarReader::new(BufReader::new(file))?;
///
/// // List available FSTs
/// for name in reader.list() {
///     println!("  {}", name);
/// }
///
/// // Read specific FST
/// if let Some(fst) = reader.read("my_fst")? {
///     println!("Loaded {} states", fst.num_states());
/// }
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
#[derive(Debug)]
pub struct FarReader<R: Read + Seek> {
    reader: R,
    entries: HashMap<String, usize>, // name -> offset
}

impl<R: Read + Seek> FarReader<R> {
    /// Creates a new FAR reader from the given reader.
    ///
    /// Reads the index from the end of the file to build an in-memory map of
    /// FST names to file offsets. This enables O(1) lookups by name.
    ///
    /// # Arguments
    ///
    /// * `reader` - The underlying reader (typically a buffered file reader)
    ///
    /// # Returns
    ///
    /// Returns a `FarReader` ready for reading FSTs by name.
    ///
    /// # Errors
    ///
    /// Returns [`Error::Io`](crate::Error::Io) if:
    /// - The file is too small to contain a valid index (< 8 bytes)
    /// - Seeking or reading fails
    /// - Entry names contain invalid UTF-8
    ///
    /// # Complexity
    ///
    /// - **Time:** O(n) where n is the number of entries in the archive
    /// - **Space:** O(n) for the in-memory index
    pub fn new(mut reader: R) -> Result<Self> {
        use byteorder::{LittleEndian, ReadBytesExt};

        // Seek to end to read index position
        let file_size = reader.seek(std::io::SeekFrom::End(0))?;

        if file_size < 8 {
            return Err(crate::Error::Io(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "FAR file too small to contain index",
            )));
        }

        // Read index position from end
        reader.seek(std::io::SeekFrom::End(-8))?;
        let index_pos = reader.read_u64::<LittleEndian>()?;

        // Seek to index
        reader.seek(std::io::SeekFrom::Start(index_pos))?;

        // Read number of entries
        let num_entries = reader.read_u32::<LittleEndian>()? as usize;
        let mut entries = HashMap::with_capacity(num_entries);

        // Read each entry
        for _ in 0..num_entries {
            let name_len = reader.read_u32::<LittleEndian>()? as usize;
            let mut name_bytes = vec![0u8; name_len];
            reader.read_exact(&mut name_bytes)?;
            let name = String::from_utf8(name_bytes).map_err(|e| {
                crate::Error::Io(std::io::Error::new(
                    std::io::ErrorKind::InvalidData,
                    format!("Invalid UTF-8 in FAR entry name: {}", e),
                ))
            })?;
            let offset = reader.read_u64::<LittleEndian>()? as usize;
            entries.insert(name, offset);
        }

        Ok(Self { reader, entries })
    }

    /// Lists all FST names in the archive.
    ///
    /// Returns references to all FST names stored in the archive. The order
    /// is not guaranteed (depends on HashMap iteration order).
    ///
    /// # Returns
    ///
    /// A vector of references to FST names.
    ///
    /// # Complexity
    ///
    /// - **Time:** O(n) where n is the number of entries
    /// - **Space:** O(n) for the returned vector
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use arcweight::io::open_far;
    /// let reader = open_far("archive.far")?;
    /// for name in reader.list() {
    ///     println!("Archive contains: {}", name);
    /// }
    /// # Ok::<(), Box<dyn std::error::Error>>(())
    /// ```
    pub fn list(&self) -> Vec<&String> {
        self.entries.keys().collect()
    }

    /// Reads an FST from the archive by name.
    ///
    /// Seeks to the FST's location in the archive and reads it using the
    /// OpenFST binary format.
    ///
    /// # Arguments
    ///
    /// * `name` - The name of the FST to read
    ///
    /// # Returns
    ///
    /// - `Ok(Some(fst))` if the FST was found and successfully read
    /// - `Ok(None)` if no FST with the given name exists in the archive
    ///
    /// # Errors
    ///
    /// Returns an error if seeking or reading the FST data fails.
    ///
    /// # Complexity
    ///
    /// - **Time:** O(1) seek + O(|V| + |E|) read
    /// - **Space:** O(|V| + |E|) for the returned FST
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use arcweight::io::open_far;
    /// use arcweight::Fst;
    ///
    /// let mut reader = open_far("archive.far")?;
    ///
    /// match reader.read("lexicon")? {
    ///     Some(fst) => println!("Loaded lexicon with {} states", fst.num_states()),
    ///     None => println!("Lexicon not found in archive"),
    /// }
    /// # Ok::<(), Box<dyn std::error::Error>>(())
    /// ```
    pub fn read(&mut self, name: &str) -> Result<Option<VectorFst<TropicalWeight>>> {
        // Find the offset for this FST
        let offset = match self.entries.get(name) {
            Some(&offset) => offset,
            None => return Ok(None),
        };

        // Seek to the offset
        self.reader.seek(std::io::SeekFrom::Start(offset as u64))?;

        // Read the FST using OpenFST format
        use crate::io::read_openfst;
        match read_openfst(&mut self.reader) {
            Ok(fst) => Ok(Some(fst)),
            Err(e) => Err(e),
        }
    }
}

/// Writer for FST Archive (FAR) files.
///
/// `FarWriter` enables creating FAR archives by sequentially adding FSTs.
/// The index is written when [`finish()`](FarWriter::finish) is called.
///
/// # Type Parameters
///
/// * `W` - The underlying writer type, must implement `Write + Seek`
///
/// # Examples
///
/// ```no_run
/// use arcweight::prelude::*;
/// use arcweight::io::{create_far, FarWriter};
/// use std::fs::File;
/// use std::io::BufWriter;
///
/// // Create using convenience function
/// let mut writer = create_far("archive.far")?;
///
/// // Or create from any Write + Seek type
/// let file = File::create("archive.far")?;
/// let mut writer = FarWriter::new(BufWriter::new(file));
///
/// // Add FSTs
/// let fst1 = VectorFst::<TropicalWeight>::new();
/// let fst2 = VectorFst::<TropicalWeight>::new();
/// writer.add("model_a", &fst1)?;
/// writer.add("model_b", &fst2)?;
///
/// // IMPORTANT: Must call finish() to write the index
/// writer.finish()?;
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
///
/// # Important
///
/// You **must** call [`finish()`](FarWriter::finish) after adding all FSTs.
/// Dropping the writer without calling `finish()` will result in an invalid
/// archive that cannot be read.
#[derive(Debug)]
pub struct FarWriter<W: Write + Seek> {
    writer: W,
    entries: Vec<(String, usize)>, // name -> offset
}

impl<W: Write + Seek> FarWriter<W> {
    /// Creates a new FAR writer.
    ///
    /// # Arguments
    ///
    /// * `writer` - The underlying writer (typically a buffered file writer)
    ///
    /// # Returns
    ///
    /// A new `FarWriter` ready to accept FSTs.
    pub fn new(writer: W) -> Self {
        Self {
            writer,
            entries: Vec::new(),
        }
    }

    /// Adds an FST to the archive.
    ///
    /// Writes the FST in OpenFST binary format and records its name and offset
    /// for the index.
    ///
    /// # Arguments
    ///
    /// * `name` - The name to associate with this FST (must be unique)
    /// * `fst` - The FST to add
    ///
    /// # Returns
    ///
    /// Returns `Ok(())` on success.
    ///
    /// # Errors
    ///
    /// Returns an error if writing the FST fails.
    ///
    /// # Complexity
    ///
    /// - **Time:** O(|V| + |E|) for writing the FST
    /// - **Space:** O(1) additional (streaming write)
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use arcweight::prelude::*;
    /// # use arcweight::io::create_far;
    /// let mut writer = create_far("archive.far")?;
    ///
    /// let fst = VectorFst::<TropicalWeight>::new();
    /// writer.add("my_fst", &fst)?;
    /// writer.finish()?;
    /// # Ok::<(), Box<dyn std::error::Error>>(())
    /// ```
    pub fn add(&mut self, name: &str, fst: &VectorFst<TropicalWeight>) -> Result<()> {
        // Write the FST in OpenFST format and record the offset
        use crate::io::write_openfst;
        let pos = self.writer.stream_position()?;
        write_openfst(fst, &mut self.writer)?;
        self.entries.push((name.to_string(), pos as usize));
        Ok(())
    }

    /// Finishes writing the archive.
    ///
    /// Writes the index containing all entry names and offsets, then flushes
    /// the writer. This method **must** be called after adding all FSTs;
    /// otherwise the archive will be invalid.
    ///
    /// # Returns
    ///
    /// Returns `Ok(())` on success. Consumes the writer.
    ///
    /// # Errors
    ///
    /// Returns an error if writing the index or flushing fails.
    ///
    /// # Complexity
    ///
    /// - **Time:** O(n) where n is the number of entries
    /// - **Space:** O(1)
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use arcweight::prelude::*;
    /// # use arcweight::io::create_far;
    /// let mut writer = create_far("archive.far")?;
    /// // ... add FSTs ...
    /// writer.finish()?; // Archive is now complete and readable
    /// # Ok::<(), Box<dyn std::error::Error>>(())
    /// ```
    pub fn finish(mut self) -> Result<()> {
        // Write index at the end
        let index_pos = self.writer.stream_position()?;

        // Write number of entries
        self.writer
            .write_u32::<LittleEndian>(self.entries.len() as u32)?;

        // Write each entry: name length, name, offset
        for (name, offset) in &self.entries {
            let name_bytes = name.as_bytes();
            self.writer
                .write_u32::<LittleEndian>(name_bytes.len() as u32)?;
            self.writer.write_all(name_bytes)?;
            self.writer.write_u64::<LittleEndian>(*offset as u64)?;
        }

        // Write index position at the very end
        self.writer.write_u64::<LittleEndian>(index_pos)?;

        self.writer.flush()?;
        Ok(())
    }
}

/// Opens a FAR file for reading.
///
/// Convenience function that opens a file and creates a buffered `FarReader`.
///
/// # Arguments
///
/// * `path` - Path to the FAR file
///
/// # Returns
///
/// Returns a `FarReader` ready to read FSTs from the archive.
///
/// # Errors
///
/// Returns an error if the file cannot be opened or the archive is invalid.
///
/// # Examples
///
/// ```no_run
/// use arcweight::io::open_far;
///
/// let mut reader = open_far("models.far")?;
/// for name in reader.list() {
///     println!("Found: {}", name);
/// }
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
pub fn open_far<P: AsRef<std::path::Path>>(
    path: P,
) -> Result<FarReader<std::io::BufReader<std::fs::File>>> {
    use std::fs::File;
    use std::io::BufReader;

    let file = File::open(path)?;
    let reader = BufReader::new(file);
    FarReader::new(reader)
}

/// Creates a FAR file for writing.
///
/// Convenience function that creates a file and returns a buffered `FarWriter`.
///
/// # Arguments
///
/// * `path` - Path to the FAR file to create (will be overwritten if exists)
///
/// # Returns
///
/// Returns a `FarWriter` ready to accept FSTs.
///
/// # Errors
///
/// Returns an error if the file cannot be created.
///
/// # Examples
///
/// ```no_run
/// use arcweight::prelude::*;
/// use arcweight::io::create_far;
///
/// let mut writer = create_far("models.far")?;
/// let fst = VectorFst::<TropicalWeight>::new();
/// writer.add("empty", &fst)?;
/// writer.finish()?;
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
pub fn create_far<P: AsRef<std::path::Path>>(
    path: P,
) -> Result<FarWriter<std::io::BufWriter<std::fs::File>>> {
    use std::io::BufWriter;
    let file = std::fs::File::create(path)?;
    let writer = BufWriter::new(file);
    Ok(FarWriter::new(writer))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::prelude::*;
    use std::io::Cursor;

    #[test]
    fn test_far_writer_new() {
        let cursor = Cursor::new(Vec::new());
        let writer = FarWriter::new(cursor);
        assert_eq!(writer.entries.len(), 0);
    }

    #[test]
    fn test_far_writer_add() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state();
        fst.set_start(s0);
        fst.set_final(s0, TropicalWeight::one());

        let cursor = Cursor::new(Vec::new());
        let mut writer = FarWriter::new(cursor);
        writer.add("test_fst", &fst).unwrap();
        assert_eq!(writer.entries.len(), 1);
        assert_eq!(writer.entries[0].0, "test_fst");
    }

    #[test]
    fn test_far_writer_finish() {
        let mut fst = VectorFst::<TropicalWeight>::new();
        let s0 = fst.add_state();
        fst.set_start(s0);
        fst.set_final(s0, TropicalWeight::one());

        let cursor = Cursor::new(Vec::new());
        let mut writer = FarWriter::new(cursor);
        writer.add("test_fst", &fst).unwrap();
        writer.finish().unwrap();
    }

    #[test]
    fn test_far_reader_list() {
        // Create a FAR file in memory
        let mut buffer = Vec::new();
        {
            let mut fst = VectorFst::<TropicalWeight>::new();
            let s0 = fst.add_state();
            fst.set_start(s0);
            fst.set_final(s0, TropicalWeight::one());

            let cursor = Cursor::new(&mut buffer);
            let mut writer = FarWriter::new(cursor);
            writer.add("fst1", &fst).unwrap();
            writer.add("fst2", &fst).unwrap();
            writer.finish().unwrap();
        }

        // Read it back
        let cursor = Cursor::new(buffer);
        let reader = FarReader::new(cursor).unwrap();
        let entries = reader.list();
        assert_eq!(entries.len(), 2);
    }
}