arcweight 0.3.0

//! Zero-copy serialization using rkyv for ultra-fast FST loading.
//!
//! This module provides extremely fast serialization and deserialization of FSTs
//! using the [rkyv](https://rkyv.org/) library, which enables zero-copy access to
//! archived data structures. This is the fastest serialization format available
//! in ArcWeight.
//!
//! # Performance Benefits
//!
//! | Metric | Traditional Deserialize | rkyv Zero-Copy |
//! |--------|-------------------------|----------------|
//! | Load time | O(n) parsing | O(1) validation |
//! | Memory allocation | Full FST copy | None (mmap) |
//! | First access | After full load | Immediate |
//! | Typical speedup | Baseline | 10-100x faster |
//!
//! # When to Use rkyv
//!
//! **Best for:**
//! - Large FSTs (millions of states)
//! - Latency-critical applications (real-time speech recognition)
//! - Frequent model loading (serverless functions)
//! - Memory-constrained environments (only accessed pages loaded)
//!
//! **Consider alternatives when:**
//! - Interoperability with OpenFST tools is required
//! - File size is critical (rkyv files are slightly larger)
//! - Cross-language compatibility is needed
//!
//! # Feature Gate
//!
//! This module requires the `zero-copy` feature:
//!
//! ```toml
//! [dependencies]
//! arcweight = { version = "0.3", features = ["zero-copy"] }
//! ```
//!
//! For memory-mapped file support, also enable `memmap2`:
//!
//! ```toml
//! [dependencies]
//! arcweight = { version = "0.3", features = ["zero-copy", "memmap2"] }
//! ```
//!
//! # Supported Weight Types
//!
//! | Weight Type | Serialize | Deserialize | Zero-Copy Access |
//! |-------------|-----------|-------------|------------------|
//! | TropicalWeight (f32) | Yes | Yes | Yes |
//! | LogWeight (f64) | Yes | Yes | Yes |
//!
//! # Examples
//!
//! ## Basic Serialization
//!
//! ```no_run
//! # #[cfg(feature = "zero-copy")]
//! # fn example() -> arcweight::Result<()> {
//! use arcweight::prelude::*;
//! use arcweight::io::rkyv_format::{serialize_tropical_fst, deserialize_tropical_fst};
//!
//! // Create an FST
//! let mut fst = VectorFst::<TropicalWeight>::new();
//! let s0 = fst.add_state();
//! let s1 = fst.add_state();
//! fst.set_start(s0);
//! fst.set_final(s1, TropicalWeight::one());
//! fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(0.5), s1));
//!
//! // Serialize to bytes
//! let bytes = serialize_tropical_fst(&fst)?;
//!
//! // Deserialize (materializes a new FST)
//! let loaded: VectorFst<TropicalWeight> = deserialize_tropical_fst(&bytes)?;
//! assert_eq!(fst.num_states(), loaded.num_states());
//! # Ok(())
//! # }
//! ```
//!
//! ## Zero-Copy Access (No Deserialization)
//!
//! ```no_run
//! # #[cfg(feature = "zero-copy")]
//! # fn example() -> arcweight::Result<()> {
//! use arcweight::prelude::*;
//! use arcweight::io::rkyv_format::{serialize_tropical_fst, access_archived_tropical_fst};
//!
//! let mut fst = VectorFst::<TropicalWeight>::new();
//! fst.add_state();
//! fst.set_start(0);
//!
//! let bytes = serialize_tropical_fst(&fst)?;
//!
//! // Access archived data directly without deserializing
//! let archived = access_archived_tropical_fst(&bytes)?;
//! println!("States: {}", archived.states.len());
//! # Ok(())
//! # }
//! ```
//!
//! ## Memory-Mapped Files
//!
//! ```no_run
//! # #[cfg(all(feature = "zero-copy", feature = "memmap2"))]
//! # fn example() -> arcweight::Result<()> {
//! use arcweight::prelude::*;
//! use arcweight::io::rkyv_format::MmapTropicalFst;
//! use std::path::Path;
//!
//! // Memory-map an FST file for instant access
//! // Only pages accessed are loaded into memory
//! let mmap_fst = MmapTropicalFst::open(Path::new("large_model.fst.rkyv"))?;
//!
//! // Access data with zero loading time
//! println!("States: {}", mmap_fst.num_states());
//! if let Some(start) = mmap_fst.start() {
//!     println!("Start state: {}", start);
//! }
//!
//! // Convert to VectorFst if needed (materializes the data)
//! let fst = mmap_fst.to_vector_fst()?;
//! # Ok(())
//! # }
//! ```
//!
//! # File Format
//!
//! The rkyv format stores FST data in a layout optimized for direct memory access:
//!
//! ```text
//! +----------------------------------+
//! | Archived States Vec              |
//! |   - Archived State 0             |
//! |     - final_weight: Option<f32>  |
//! |     - arcs: Vec<ArchivedArc>     |
//! |   - Archived State 1             |
//! |     ...                          |
//! +----------------------------------+
//! | Archived Start State: Option<u32>|
//! +----------------------------------+
//! ```
//!
//! # Safety Considerations
//!
//! When using memory-mapped files:
//! - The file must not be modified while the mapping is active
//! - The file must contain valid rkyv-serialized data
//! - The mapping is validated on open to ensure safety
//!
//! # References
//!
//! - rkyv documentation: <https://rkyv.org/>
//! - Zero-copy deserialization explained: <https://rkyv.org/zero-copy-deserialization.html>

use crate::semiring::{LogWeight, Semiring, TropicalWeight};
use crate::Error;
use crate::Result;

use rkyv::{rancor, Archive, Deserialize, Serialize};

/// Archived representation of an FST arc with f32 weight.
///
/// This struct is used for serializing arcs in [`TropicalWeight`] FSTs.
/// It stores the minimal data needed to reconstruct an arc.
///
/// # Fields
///
/// All fields are public to allow direct access from archived data.
#[derive(Archive, Deserialize, Serialize, Debug, Clone, PartialEq)]
pub struct RkyvArcF32 {
    /// Input label (symbol ID on the input tape).
    pub ilabel: u32,
    /// Output label (symbol ID on the output tape).
    pub olabel: u32,
    /// Weight value as f32 (tropical semiring uses -log probability).
    pub weight: f32,
    /// Destination state ID.
    pub nextstate: u32,
}

/// Archived representation of an FST state with f32 weights.
///
/// This struct stores all data for a single state in a [`TropicalWeight`] FST.
#[derive(Archive, Deserialize, Serialize, Debug, Clone)]
pub struct RkyvStateF32 {
    /// Final weight, or `None` if the state is non-final.
    pub final_weight: Option<f32>,
    /// Outgoing arcs from this state.
    pub arcs: Vec<RkyvArcF32>,
}

/// Archived [`TropicalWeight`] FST representation.
///
/// This is the top-level structure for rkyv-serialized FSTs using the
/// tropical semiring (f32 weights). It can be accessed zero-copy via
/// [`access_archived_tropical_fst`] or materialized to a [`VectorFst`](crate::fst::VectorFst)
/// via [`deserialize_tropical_fst`].
///
/// # Example
///
/// ```no_run
/// use arcweight::io::rkyv_format::{serialize_tropical_fst, access_archived_tropical_fst};
/// use arcweight::prelude::*;
///
/// let fst = VectorFst::<TropicalWeight>::new();
/// let bytes = serialize_tropical_fst(&fst)?;
/// let archived = access_archived_tropical_fst(&bytes)?;
///
/// // Direct access to archived data
/// println!("States: {}", archived.states.len());
/// # Ok::<(), arcweight::Error>(())
/// ```
#[derive(Archive, Deserialize, Serialize, Debug, Clone)]
pub struct RkyvFstF32 {
    /// All states in the FST, indexed by state ID.
    pub states: Vec<RkyvStateF32>,
    /// Start state ID, or `None` if the FST has no start state.
    pub start: Option<u32>,
}

/// Archived representation of an FST arc with f64 weight.
///
/// This struct is used for serializing arcs in [`LogWeight`] FSTs.
/// It stores the minimal data needed to reconstruct an arc.
#[derive(Archive, Deserialize, Serialize, Debug, Clone, PartialEq)]
pub struct RkyvArcF64 {
    /// Input label (symbol ID on the input tape).
    pub ilabel: u32,
    /// Output label (symbol ID on the output tape).
    pub olabel: u32,
    /// Weight value as f64 (log semiring uses -log probability).
    pub weight: f64,
    /// Destination state ID.
    pub nextstate: u32,
}

/// Archived representation of an FST state with f64 weights.
///
/// This struct stores all data for a single state in a [`LogWeight`] FST.
#[derive(Archive, Deserialize, Serialize, Debug, Clone)]
pub struct RkyvStateF64 {
    /// Final weight, or `None` if the state is non-final.
    pub final_weight: Option<f64>,
    /// Outgoing arcs from this state.
    pub arcs: Vec<RkyvArcF64>,
}

/// Archived [`LogWeight`] FST representation.
///
/// This is the top-level structure for rkyv-serialized FSTs using the
/// log semiring (f64 weights). It can be accessed zero-copy via
/// [`access_archived_log_fst`] or materialized to a [`VectorFst`](crate::fst::VectorFst)
/// via [`deserialize_log_fst`].
#[derive(Archive, Deserialize, Serialize, Debug, Clone)]
pub struct RkyvFstF64 {
    /// All states in the FST, indexed by state ID.
    pub states: Vec<RkyvStateF64>,
    /// Start state ID, or `None` if the FST has no start state.
    pub start: Option<u32>,
}

/// Serializes a [`TropicalWeight`] FST to bytes using rkyv.
///
/// Produces a compact binary representation optimized for zero-copy access.
/// The resulting bytes can be written to a file, stored in memory, or
/// memory-mapped for instant access to large FSTs.
///
/// # Type Parameters
///
/// * `F` - The FST type, must implement `Fst<TropicalWeight>`
///
/// # Arguments
///
/// * `fst` - Reference to the FST to serialize
///
/// # Returns
///
/// Returns the serialized bytes on success.
///
/// # Errors
///
/// Returns [`Error::Serialization`] if rkyv
/// serialization fails.
///
/// # Complexity
///
/// - **Time:** O(|V| + |E|) where |V| is states and |E| is arcs
/// - **Space:** O(|V| + |E|) for the output bytes
///
/// # Examples
///
/// ```no_run
/// # #[cfg(feature = "zero-copy")]
/// # fn example() -> arcweight::Result<()> {
/// use arcweight::prelude::*;
/// use arcweight::io::rkyv_format::serialize_tropical_fst;
///
/// let fst = VectorFst::<TropicalWeight>::new();
/// let bytes = serialize_tropical_fst(&fst)?;
/// println!("Serialized to {} bytes", bytes.len());
/// # Ok(())
/// # }
/// ```
pub fn serialize_tropical_fst<F>(fst: &F) -> Result<Vec<u8>>
where
    F: crate::fst::Fst<TropicalWeight>,
{
    let rkyv_fst = to_rkyv_fst_tropical(fst);

    let bytes = rkyv::to_bytes::<rancor::Error>(&rkyv_fst)
        .map_err(|e| Error::Serialization(format!("rkyv serialization failed: {}", e)))?;

    Ok(bytes.to_vec())
}

/// Deserializes a [`TropicalWeight`] FST from rkyv bytes.
///
/// Materializes a complete [`VectorFst`](crate::fst::VectorFst) from archived data. For zero-copy
/// access without materialization, use [`access_archived_tropical_fst`] instead.
///
/// # Arguments
///
/// * `bytes` - The rkyv-serialized bytes
///
/// # Returns
///
/// Returns the deserialized FST on success.
///
/// # Errors
///
/// Returns [`Error::Serialization`] if the bytes
/// are not valid rkyv data or do not represent a valid FST.
///
/// # Complexity
///
/// - **Time:** O(|V| + |E|) to construct the FST
/// - **Space:** O(|V| + |E|) for the output FST
///
/// # Examples
///
/// ```no_run
/// # #[cfg(feature = "zero-copy")]
/// # fn example() -> arcweight::Result<()> {
/// use arcweight::prelude::*;
/// use arcweight::io::rkyv_format::{serialize_tropical_fst, deserialize_tropical_fst};
///
/// let original = VectorFst::<TropicalWeight>::new();
/// let bytes = serialize_tropical_fst(&original)?;
/// let loaded = deserialize_tropical_fst(&bytes)?;
/// assert_eq!(original.num_states(), loaded.num_states());
/// # Ok(())
/// # }
/// ```
pub fn deserialize_tropical_fst(bytes: &[u8]) -> Result<crate::fst::VectorFst<TropicalWeight>> {
    let archived = access_archived_tropical_fst(bytes)?;
    from_archived_fst_tropical(archived)
}

/// Accesses archived [`TropicalWeight`] FST data without deserialization.
///
/// This is the zero-copy access function. It validates and returns a reference
/// to the archived FST data, which can be traversed directly without allocating
/// a new FST structure.
///
/// # Arguments
///
/// * `bytes` - The rkyv-serialized bytes
///
/// # Returns
///
/// Returns a reference to the archived FST data. The returned reference has
/// the same lifetime as the input bytes.
///
/// # Errors
///
/// Returns [`Error::Serialization`] if validation
/// fails (bytes are not valid rkyv data).
///
/// # Complexity
///
/// - **Time:** O(1) for validation
/// - **Space:** O(1) - no allocation
///
/// # Examples
///
/// ```no_run
/// use arcweight::prelude::*;
/// use arcweight::io::rkyv_format::{serialize_tropical_fst, access_archived_tropical_fst};
///
/// let fst = VectorFst::<TropicalWeight>::new();
/// let bytes = serialize_tropical_fst(&fst)?;
///
/// // Zero-copy access - no FST allocation
/// let archived = access_archived_tropical_fst(&bytes)?;
/// println!("States: {}", archived.states.len());
/// # Ok::<(), arcweight::Error>(())
/// ```
pub fn access_archived_tropical_fst(bytes: &[u8]) -> Result<&ArchivedRkyvFstF32> {
    rkyv::access::<ArchivedRkyvFstF32, rancor::Error>(bytes)
        .map_err(|e| Error::Serialization(format!("rkyv validation failed: {}", e)))
}

/// Serializes a [`LogWeight`] FST to bytes using rkyv.
///
/// Produces a compact binary representation optimized for zero-copy access.
/// This is the f64-weight variant for log semiring FSTs.
///
/// # Type Parameters
///
/// * `F` - The FST type, must implement `Fst<LogWeight>`
///
/// # Arguments
///
/// * `fst` - Reference to the FST to serialize
///
/// # Returns
///
/// Returns the serialized bytes on success.
///
/// # Errors
///
/// Returns [`Error::Serialization`] if rkyv
/// serialization fails.
pub fn serialize_log_fst<F>(fst: &F) -> Result<Vec<u8>>
where
    F: crate::fst::Fst<LogWeight>,
{
    let rkyv_fst = to_rkyv_fst_log(fst);

    let bytes = rkyv::to_bytes::<rancor::Error>(&rkyv_fst)
        .map_err(|e| Error::Serialization(format!("rkyv serialization failed: {}", e)))?;

    Ok(bytes.to_vec())
}

/// Deserializes a [`LogWeight`] FST from rkyv bytes.
///
/// Materializes a complete [`VectorFst`](crate::fst::VectorFst) from archived data. This is the
/// f64-weight variant for log semiring FSTs.
///
/// # Arguments
///
/// * `bytes` - The rkyv-serialized bytes
///
/// # Returns
///
/// Returns the deserialized FST on success.
///
/// # Errors
///
/// Returns [`Error::Serialization`] if the bytes
/// are not valid rkyv data.
pub fn deserialize_log_fst(bytes: &[u8]) -> Result<crate::fst::VectorFst<LogWeight>> {
    let archived = access_archived_log_fst(bytes)?;
    from_archived_fst_log(archived)
}

/// Accesses archived [`LogWeight`] FST data without deserialization.
///
/// This is the zero-copy access function for log semiring FSTs.
///
/// # Arguments
///
/// * `bytes` - The rkyv-serialized bytes
///
/// # Returns
///
/// Returns a reference to the archived FST data.
///
/// # Errors
///
/// Returns [`Error::Serialization`] if validation fails.
pub fn access_archived_log_fst(bytes: &[u8]) -> Result<&ArchivedRkyvFstF64> {
    rkyv::access::<ArchivedRkyvFstF64, rancor::Error>(bytes)
        .map_err(|e| Error::Serialization(format!("rkyv validation failed: {}", e)))
}

/// Convert a TropicalWeight FST to RkyvFstF32 representation
fn to_rkyv_fst_tropical<F>(fst: &F) -> RkyvFstF32
where
    F: crate::fst::Fst<TropicalWeight>,
{
    let mut states = Vec::with_capacity(fst.num_states());

    for state_id in 0..fst.num_states() as u32 {
        let final_weight = fst.final_weight(state_id).map(|w| *w.value());
        let arcs: Vec<RkyvArcF32> = fst
            .arcs(state_id)
            .map(|arc| RkyvArcF32 {
                ilabel: arc.ilabel,
                olabel: arc.olabel,
                weight: *arc.weight.value(),
                nextstate: arc.nextstate,
            })
            .collect();

        states.push(RkyvStateF32 { final_weight, arcs });
    }

    RkyvFstF32 {
        states,
        start: fst.start(),
    }
}

/// Convert a LogWeight FST to RkyvFstF64 representation
fn to_rkyv_fst_log<F>(fst: &F) -> RkyvFstF64
where
    F: crate::fst::Fst<LogWeight>,
{
    let mut states = Vec::with_capacity(fst.num_states());

    for state_id in 0..fst.num_states() as u32 {
        let final_weight = fst.final_weight(state_id).map(|w| *w.value());
        let arcs: Vec<RkyvArcF64> = fst
            .arcs(state_id)
            .map(|arc| RkyvArcF64 {
                ilabel: arc.ilabel,
                olabel: arc.olabel,
                weight: *arc.weight.value(),
                nextstate: arc.nextstate,
            })
            .collect();

        states.push(RkyvStateF64 { final_weight, arcs });
    }

    RkyvFstF64 {
        states,
        start: fst.start(),
    }
}

/// Convert an archived FST to TropicalWeight VectorFst
fn from_archived_fst_tropical(
    archived: &ArchivedRkyvFstF32,
) -> Result<crate::fst::VectorFst<TropicalWeight>> {
    use crate::arc::Arc;
    use crate::fst::MutableFst;
    use rkyv::option::ArchivedOption;

    let mut fst = crate::fst::VectorFst::new();

    // Add all states
    for _ in 0..archived.states.len() {
        fst.add_state();
    }

    // Set start state
    if let ArchivedOption::Some(start) = &archived.start {
        let start_val: u32 = (*start).into();
        fst.set_start(start_val);
    }

    // Add arcs and final weights
    for (state_id, archived_state) in archived.states.iter().enumerate() {
        let state_id = state_id as u32;

        if let ArchivedOption::Some(weight_val) = &archived_state.final_weight {
            let val: f32 = (*weight_val).into();
            fst.set_final(state_id, TropicalWeight::new(val));
        }

        for archived_arc in archived_state.arcs.iter() {
            let ilabel: u32 = archived_arc.ilabel.into();
            let olabel: u32 = archived_arc.olabel.into();
            let weight: f32 = archived_arc.weight.into();
            let nextstate: u32 = archived_arc.nextstate.into();

            fst.add_arc(
                state_id,
                Arc::new(ilabel, olabel, TropicalWeight::new(weight), nextstate),
            );
        }
    }

    Ok(fst)
}

/// Convert an archived FST to LogWeight VectorFst
fn from_archived_fst_log(
    archived: &ArchivedRkyvFstF64,
) -> Result<crate::fst::VectorFst<LogWeight>> {
    use crate::arc::Arc;
    use crate::fst::MutableFst;
    use rkyv::option::ArchivedOption;

    let mut fst = crate::fst::VectorFst::new();

    for _ in 0..archived.states.len() {
        fst.add_state();
    }

    if let ArchivedOption::Some(start) = &archived.start {
        let start_val: u32 = (*start).into();
        fst.set_start(start_val);
    }

    for (state_id, archived_state) in archived.states.iter().enumerate() {
        let state_id = state_id as u32;

        if let ArchivedOption::Some(weight_val) = &archived_state.final_weight {
            let val: f64 = (*weight_val).into();
            fst.set_final(state_id, LogWeight::new(val));
        }

        for archived_arc in archived_state.arcs.iter() {
            let ilabel: u32 = archived_arc.ilabel.into();
            let olabel: u32 = archived_arc.olabel.into();
            let weight: f64 = archived_arc.weight.into();
            let nextstate: u32 = archived_arc.nextstate.into();

            fst.add_arc(
                state_id,
                Arc::new(ilabel, olabel, LogWeight::new(weight), nextstate),
            );
        }
    }

    Ok(fst)
}

/// Writes a [`TropicalWeight`] FST to a file in rkyv format.
///
/// Convenience function that serializes an FST and writes it to a file.
///
/// # Arguments
///
/// * `fst` - Reference to the FST to serialize
/// * `path` - Path to the output file (will be created/overwritten)
///
/// # Returns
///
/// Returns `Ok(())` on success.
///
/// # Errors
///
/// Returns an error if file creation or writing fails.
///
/// # Examples
///
/// ```no_run
/// # #[cfg(feature = "zero-copy")]
/// # fn example() -> arcweight::Result<()> {
/// use arcweight::prelude::*;
/// use arcweight::io::rkyv_format::write_tropical_rkyv;
/// use std::path::Path;
///
/// let fst = VectorFst::<TropicalWeight>::new();
/// write_tropical_rkyv(&fst, Path::new("model.fst.rkyv"))?;
/// # Ok(())
/// # }
/// ```
pub fn write_tropical_rkyv<F>(fst: &F, path: &std::path::Path) -> Result<()>
where
    F: crate::fst::Fst<TropicalWeight>,
{
    use std::io::Write;
    let bytes = serialize_tropical_fst(fst)?;
    let mut file = std::fs::File::create(path)?;
    file.write_all(&bytes)?;
    Ok(())
}

/// Reads a [`TropicalWeight`] FST from a file in rkyv format.
///
/// Convenience function that reads a file and deserializes it to an FST.
/// For zero-copy access to large files, use [`MmapTropicalFst`] instead.
///
/// # Arguments
///
/// * `path` - Path to the rkyv file
///
/// # Returns
///
/// Returns the deserialized FST on success.
///
/// # Errors
///
/// Returns an error if file reading or deserialization fails.
///
/// # Examples
///
/// ```no_run
/// # #[cfg(feature = "zero-copy")]
/// # fn example() -> arcweight::Result<()> {
/// use arcweight::prelude::*;
/// use arcweight::io::rkyv_format::read_tropical_rkyv;
/// use std::path::Path;
///
/// let fst = read_tropical_rkyv(Path::new("model.fst.rkyv"))?;
/// println!("Loaded {} states", fst.num_states());
/// # Ok(())
/// # }
/// ```
pub fn read_tropical_rkyv(path: &std::path::Path) -> Result<crate::fst::VectorFst<TropicalWeight>> {
    let bytes = std::fs::read(path)?;
    deserialize_tropical_fst(&bytes)
}

/// Writes a [`LogWeight`] FST to a file in rkyv format.
///
/// Convenience function that serializes an FST and writes it to a file.
/// This is the f64-weight variant for log semiring FSTs.
///
/// # Arguments
///
/// * `fst` - Reference to the FST to serialize
/// * `path` - Path to the output file (will be created/overwritten)
///
/// # Returns
///
/// Returns `Ok(())` on success.
///
/// # Errors
///
/// Returns an error if file creation or writing fails.
pub fn write_log_rkyv<F>(fst: &F, path: &std::path::Path) -> Result<()>
where
    F: crate::fst::Fst<LogWeight>,
{
    use std::io::Write;
    let bytes = serialize_log_fst(fst)?;
    let mut file = std::fs::File::create(path)?;
    file.write_all(&bytes)?;
    Ok(())
}

/// Reads a [`LogWeight`] FST from a file in rkyv format.
///
/// Convenience function that reads a file and deserializes it to an FST.
/// For zero-copy access to large files, use [`MmapLogFst`] instead.
///
/// # Arguments
///
/// * `path` - Path to the rkyv file
///
/// # Returns
///
/// Returns the deserialized FST on success.
///
/// # Errors
///
/// Returns an error if file reading or deserialization fails.
pub fn read_log_rkyv(path: &std::path::Path) -> Result<crate::fst::VectorFst<LogWeight>> {
    let bytes = std::fs::read(path)?;
    deserialize_log_fst(&bytes)
}

// Memory-mapped FST support
#[cfg(feature = "memmap2")]
mod mmap_support {
    use super::*;
    use memmap2::Mmap;
    use std::fs::File;
    use std::path::Path;

    /// Memory-mapped [`TropicalWeight`] FST for instant zero-copy access.
    ///
    /// `MmapTropicalFst` uses memory mapping to provide instant access to FST data
    /// stored in files. The operating system pages in data on demand, making this
    /// ideal for large FSTs that may not fit entirely in memory.
    ///
    /// # Performance Characteristics
    ///
    /// | Operation | Complexity | Notes |
    /// |-----------|------------|-------|
    /// | Open | O(1) | Just sets up memory mapping |
    /// | num_states | O(1) | Direct archived access |
    /// | start | O(1) | Direct archived access |
    /// | to_vector_fst | O(\|V\| + \|E\|) | Materializes full FST |
    ///
    /// # Memory Behavior
    ///
    /// - Only pages accessed are loaded into physical memory
    /// - Pages may be evicted under memory pressure
    /// - Total virtual address space = file size
    /// - Ideal for large FSTs with sparse access patterns
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # #[cfg(all(feature = "zero-copy", feature = "memmap2"))]
    /// # fn example() -> arcweight::Result<()> {
    /// use arcweight::io::rkyv_format::MmapTropicalFst;
    /// use std::path::Path;
    ///
    /// // Instant loading regardless of FST size
    /// let fst = MmapTropicalFst::open(Path::new("huge_model.fst.rkyv"))?;
    ///
    /// // Access metadata without loading full FST
    /// println!("States: {}", fst.num_states());
    ///
    /// // Convert to VectorFst when needed (loads all data)
    /// let vector_fst = fst.to_vector_fst()?;
    /// # Ok(())
    /// # }
    /// ```
    pub struct MmapTropicalFst {
        mmap: Mmap,
    }

    impl MmapTropicalFst {
        /// Opens a memory-mapped FST file.
        ///
        /// Creates a memory mapping for the file and validates that it contains
        /// valid rkyv-serialized FST data.
        ///
        /// # Arguments
        ///
        /// * `path` - Path to the rkyv-format FST file
        ///
        /// # Returns
        ///
        /// Returns a `MmapTropicalFst` on success.
        ///
        /// # Errors
        ///
        /// Returns an error if:
        /// - The file cannot be opened
        /// - Memory mapping fails
        /// - The file does not contain valid rkyv FST data
        ///
        /// # Safety
        ///
        /// The file must contain valid rkyv-serialized FST data. The file should
        /// not be modified while the `MmapTropicalFst` is in use. This function
        /// uses `unsafe` internally to create the memory mapping.
        pub fn open(path: &Path) -> Result<Self> {
            let file = File::open(path)?;
            let mmap = unsafe { Mmap::map(&file) }.map_err(Error::Io)?;

            // Validate the archive
            let _ = access_archived_tropical_fst(&mmap)?;

            Ok(Self { mmap })
        }

        /// Returns a reference to the archived FST data.
        ///
        /// Provides direct access to the archived data structure for advanced
        /// use cases. The returned reference can be used to traverse the FST
        /// structure without materialization.
        pub fn archived(&self) -> &ArchivedRkyvFstF32 {
            rkyv::access::<ArchivedRkyvFstF32, rancor::Error>(&self.mmap)
                .expect("Archive was validated during open")
        }

        /// Returns the number of states in the FST.
        pub fn num_states(&self) -> usize {
            self.archived().states.len()
        }

        /// Returns the start state, or `None` if the FST has no start state.
        pub fn start(&self) -> Option<u32> {
            use rkyv::option::ArchivedOption;
            match &self.archived().start {
                ArchivedOption::Some(s) => Some((*s).into()),
                ArchivedOption::None => None,
            }
        }

        /// Returns the raw bytes of the memory mapping.
        ///
        /// Useful for passing to other functions that accept byte slices.
        pub fn as_bytes(&self) -> &[u8] {
            &self.mmap
        }

        /// Converts to a full [`VectorFst`](crate::fst::VectorFst), materializing the archived data.
        ///
        /// This reads all data from the memory mapping and constructs a
        /// complete FST in memory. Use this when you need to modify the FST
        /// or use it with algorithms that require `MutableFst`.
        ///
        /// # Returns
        ///
        /// Returns the materialized FST on success.
        ///
        /// # Complexity
        ///
        /// - **Time:** O(|V| + |E|)
        /// - **Space:** O(|V| + |E|) for the new FST
        pub fn to_vector_fst(&self) -> Result<crate::fst::VectorFst<TropicalWeight>> {
            from_archived_fst_tropical(self.archived())
        }
    }

    impl std::fmt::Debug for MmapTropicalFst {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            f.debug_struct("MmapTropicalFst")
                .field("num_states", &self.num_states())
                .field("start", &self.start())
                .finish()
        }
    }

    /// Memory-mapped [`LogWeight`] FST for instant zero-copy access.
    ///
    /// This is the f64-weight variant of [`MmapTropicalFst`] for log semiring FSTs.
    /// See [`MmapTropicalFst`] for detailed documentation on memory-mapped FST usage.
    pub struct MmapLogFst {
        mmap: Mmap,
    }

    impl MmapLogFst {
        /// Opens a memory-mapped FST file.
        ///
        /// See [`MmapTropicalFst::open`] for detailed documentation.
        pub fn open(path: &Path) -> Result<Self> {
            let file = File::open(path)?;
            let mmap = unsafe { Mmap::map(&file) }.map_err(Error::Io)?;

            let _ = access_archived_log_fst(&mmap)?;

            Ok(Self { mmap })
        }

        /// Returns a reference to the archived FST data.
        pub fn archived(&self) -> &ArchivedRkyvFstF64 {
            rkyv::access::<ArchivedRkyvFstF64, rancor::Error>(&self.mmap)
                .expect("Archive was validated during open")
        }

        /// Returns the number of states in the FST.
        pub fn num_states(&self) -> usize {
            self.archived().states.len()
        }

        /// Returns the start state, or `None` if the FST has no start state.
        pub fn start(&self) -> Option<u32> {
            use rkyv::option::ArchivedOption;
            match &self.archived().start {
                ArchivedOption::Some(s) => Some((*s).into()),
                ArchivedOption::None => None,
            }
        }

        /// Converts to a full [`VectorFst`](crate::fst::VectorFst), materializing the archived data.
        pub fn to_vector_fst(&self) -> Result<crate::fst::VectorFst<LogWeight>> {
            from_archived_fst_log(self.archived())
        }
    }

    impl std::fmt::Debug for MmapLogFst {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            f.debug_struct("MmapLogFst")
                .field("num_states", &self.num_states())
                .field("start", &self.start())
                .finish()
        }
    }
}

#[cfg(feature = "memmap2")]
pub use mmap_support::{MmapLogFst, MmapTropicalFst};

#[cfg(test)]
mod tests {
    use super::*;
    use crate::fst::MutableFst;
    use crate::prelude::*;

    fn create_test_fst() -> VectorFst<TropicalWeight> {
        let mut fst = VectorFst::new();
        let s0 = fst.add_state();
        let s1 = fst.add_state();
        let s2 = fst.add_state();

        fst.set_start(s0);
        fst.set_final(s2, TropicalWeight::new(0.5));

        fst.add_arc(s0, crate::arc::Arc::new(1, 2, TropicalWeight::new(1.0), s1));
        fst.add_arc(s1, crate::arc::Arc::new(3, 4, TropicalWeight::new(2.0), s2));
        fst.add_arc(s0, crate::arc::Arc::new(5, 5, TropicalWeight::new(0.5), s2));

        fst
    }

    #[test]
    fn test_tropical_rkyv_roundtrip() {
        let original = create_test_fst();

        // Serialize
        let bytes = serialize_tropical_fst(&original).expect("Serialization failed");

        // Deserialize
        let loaded = deserialize_tropical_fst(&bytes).expect("Deserialization failed");

        // Verify
        assert_eq!(original.num_states(), loaded.num_states());
        assert_eq!(original.start(), loaded.start());

        for state in 0..original.num_states() as u32 {
            assert_eq!(original.final_weight(state), loaded.final_weight(state));

            let orig_arcs: Vec<_> = original.arcs(state).collect();
            let loaded_arcs: Vec<_> = loaded.arcs(state).collect();
            assert_eq!(orig_arcs, loaded_arcs);
        }
    }

    #[test]
    fn test_zero_copy_access() {
        let fst = create_test_fst();
        let bytes = serialize_tropical_fst(&fst).expect("Serialization failed");

        // Access without deserialization
        let archived = access_archived_tropical_fst(&bytes).expect("Access failed");

        // Verify archived data
        assert_eq!(archived.states.len(), 3);
        assert!(archived.start.is_some());
    }

    #[test]
    fn test_empty_fst() {
        let fst = VectorFst::<TropicalWeight>::new();
        let bytes = serialize_tropical_fst(&fst).expect("Serialization failed");
        let loaded = deserialize_tropical_fst(&bytes).expect("Deserialization failed");

        assert_eq!(loaded.num_states(), 0);
        assert_eq!(loaded.start(), None);
    }

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

        let bytes = serialize_tropical_fst(&fst).expect("Serialization failed");
        let loaded = deserialize_tropical_fst(&bytes).expect("Deserialization failed");

        assert_eq!(loaded.num_states(), 1);
        assert_eq!(loaded.start(), Some(0));
        assert_eq!(loaded.final_weight(0), Some(&TropicalWeight::one()));
    }

    #[test]
    fn test_log_weight_roundtrip() {
        let mut fst = VectorFst::<LogWeight>::new();
        let s0 = fst.add_state();
        let s1 = fst.add_state();
        fst.set_start(s0);
        fst.set_final(s1, LogWeight::new(0.5));
        fst.add_arc(s0, crate::arc::Arc::new(1, 1, LogWeight::new(1.0), s1));

        let bytes = serialize_log_fst(&fst).expect("Serialization failed");
        let loaded = deserialize_log_fst(&bytes).expect("Deserialization failed");

        assert_eq!(fst.num_states(), loaded.num_states());
        assert_eq!(fst.start(), loaded.start());
    }
}