//! I/O operations for TACO format
//!
//! This module provides high-level readers and writers for TACO files,
//! handling chunking, indexing, and frame access.

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use std::fs::File;
use std::io::{BufReader, BufWriter, Seek, SeekFrom, Write};
use std::path::Path;

use crate::compression::{CompressionSettings, PrecisionMode};
use crate::frame::{Frame, FrameData};
use crate::metadata::{AtomMetadata, Constraint, SimulationMetadata};
use crate::{Error, Header, Result};

/// Writer for TACO trajectory files
pub struct Writer {
    /// The file writer
    writer: BufWriter<File>,

    /// Header information
    header: Header,

    /// Current number of frames written
    num_frames: u64,

    /// Offsets to frames for the index table
    frame_offsets: Vec<u64>,

    /// Compression settings
    compression_settings: CompressionSettings,

    /// Reference frame for delta encoding
    reference_frame: Option<FrameData>,

    /// Has the header been written to disk yet?
    header_written: bool,

    /// Maximum reserved length for the header JSON
    header_max_length: usize,
}

impl Writer {
    /// Create a new TACO writer
    pub fn create<P: AsRef<Path>>(
        path: P,
        num_atoms: u32,
        time_step: f64,
        simulation_metadata: SimulationMetadata,
        atom_metadata: AtomMetadata,
        compression_settings: CompressionSettings,
    ) -> Result<Self> {
        // Validate atom metadata
        if !atom_metadata.validate(num_atoms as usize) {
            return Err(Error::InvalidMetadata(
                "Atom metadata has inconsistent lengths".to_string(),
            ));
        }

        // Create the file
        let file = File::create(path)?;
        // Use a larger buffer (1MB)
        let writer = BufWriter::with_capacity(1024 * 1024, file);

        // Create the header
        let header = Header::new(
            num_atoms,
            time_step,
            simulation_metadata,
            atom_metadata,
            compression_settings.clone(),
        );

        Ok(Writer {
            writer,
            header,
            num_frames: 0,
            frame_offsets: Vec::new(),
            compression_settings,
            reference_frame: None,
            header_written: false,
            header_max_length: 0,
        })
    }

    /// Open an existing TACO file for appending frames
    ///
    /// This method allows you to add more frames to an existing TACO trajectory file
    /// without rewriting the entire file. The append operation:
    ///
    /// - Reads the existing file header and frame index
    /// - Loads the last frame as the reference for delta encoding
    /// - Positions the file pointer at the end of the frame data
    /// - Allows writing new frames with proper delta encoding chain
    /// - Updates the frame index and header when `finish()` is called
    ///
    /// # Important Notes
    ///
    /// - The delta encoding chain is maintained by using the last existing frame
    ///   as the reference for the first appended frame
    /// - Full frame intervals are respected based on the total frame count
    /// - All metadata (atom count, compression settings, etc.) is preserved
    /// - The original frame index is overwritten when new frames are appended
    ///
    /// # Example
    ///
    /// ```no_run
    /// use taco_format::{Writer, Reader, Frame, FrameData, SimulationMetadata, AtomMetadata, CompressionSettings};
    /// use ndarray::Array2;
    ///
    /// # fn example() -> taco_format::Result<()> {
    /// // Create an initial trajectory file
    /// let mut writer = Writer::create(
    ///     "trajectory.taco",
    ///     3,
    ///     0.001,
    ///     SimulationMetadata::default(),
    ///     AtomMetadata::default(),
    ///     CompressionSettings::default(),
    /// )?;
    ///
    /// // Write some initial frames
    /// let positions = Array2::<f32>::zeros((3, 3));
    /// for i in 0..10 {
    ///     let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
    ///     writer.write_frame(frame)?;
    /// }
    /// writer.finish()?;
    ///
    /// // Later, append more frames
    /// let mut writer = Writer::append("trajectory.taco")?;
    /// for i in 10..20 {
    ///     let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
    ///     writer.write_frame(frame)?;
    /// }
    /// writer.finish()?;
    ///
    /// // Verify we now have 20 frames
    /// let reader = Reader::open("trajectory.taco")?;
    /// assert_eq!(reader.num_frames(), 20);
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The file does not exist or cannot be opened
    /// - The file is not a valid TACO format file
    /// - The file is corrupted or has an invalid frame index
    pub fn append<P: AsRef<Path>>(path: P) -> Result<Self> {
        // First, open the file for reading to get existing data
        let mut reader = Reader::open(&path)?;
        let header = reader.header().clone();
        let num_frames = header.num_frames;
        let frame_offsets = reader.frame_offsets.clone();
        let compression_settings = header.compression_settings.clone();
        let header_max_length = reader.header_length;

        // Read the last frame to use as reference for delta encoding
        let reference_frame = if num_frames > 0 {
            let last_frame = reader.read_frame(num_frames - 1)?;
            Some(last_frame.data)
        } else {
            None
        };

        // Get the position where the frame index starts (we'll overwrite it)
        let frame_index_offset = header.frame_index_offset;

        // Close the reader by dropping it
        drop(reader);

        // Open the file for writing in append mode, then seek to frame index position
        let file = std::fs::OpenOptions::new()
            .read(true)
            .write(true)
            .open(&path)?;

        let mut writer = BufWriter::with_capacity(1024 * 1024, file);

        // Seek to where the frame index table starts (we'll overwrite it with new frames)
        writer.seek(SeekFrom::Start(frame_index_offset))?;

        Ok(Writer {
            writer,
            header,
            num_frames,
            frame_offsets,
            compression_settings,
            reference_frame,
            header_written: true,
            header_max_length,
        })
    }

    /// Create a new TACO writer with a custom header
    pub fn create_with_header<P: AsRef<Path>>(path: P, header: Header) -> Result<Self> {
        // Validate atom metadata
        if !header.atom_metadata.validate(header.num_atoms as usize) {
            return Err(Error::InvalidMetadata(
                "Atom metadata has inconsistent lengths".to_string(),
            ));
        }

        let num_atoms = header.num_atoms as usize;
        for constraint in &header.constraints {
            if let Constraint::FixAtoms { indices } = constraint {
                if let Some(&bad_index) = indices.iter().find(|&&i| i >= num_atoms) {
                    return Err(Error::InvalidMetadata(format!(
                        "Constraint FixAtoms index {bad_index} out of bounds for {num_atoms} atoms"
                    )));
                }
            }
        }

        // Create the file
        let file = File::create(path)?;
        // Use a larger buffer (1MB)
        let writer = BufWriter::with_capacity(1024 * 1024, file);

        // Clone needed settings before writing header
        let compression_settings = header.compression_settings.clone();

        Ok(Writer {
            writer,
            header,
            num_frames: 0,
            frame_offsets: Vec::new(),
            compression_settings,
            reference_frame: None,
            header_written: false,
            header_max_length: 0,
        })
    }

    /// Set constraints for the header (call before first write_frame)
    pub fn set_constraints(&mut self, constraints: Vec<Constraint>) {
        self.header.constraints = constraints;
    }

    /// Set periodic boundary conditions for the header (call before first write_frame)
    pub fn set_pbc(&mut self, pbc: [bool; 3]) {
        self.header.pbc = pbc;
    }

    fn validate_constraints(&self) -> Result<()> {
        let num_atoms = self.header.num_atoms as usize;

        for constraint in &self.header.constraints {
            if let Constraint::FixAtoms { indices } = constraint {
                if let Some(&bad_index) = indices.iter().find(|&&i| i >= num_atoms) {
                    return Err(Error::InvalidMetadata(format!(
                        "Constraint FixAtoms index {bad_index} out of bounds for {num_atoms} atoms"
                    )));
                }
            }
        }

        Ok(())
    }

    /// Ensure the header is written to disk before writing the first frame.
    fn ensure_header_written(&mut self) -> Result<()> {
        if !self.header_written {
            self.validate_constraints()?;

            // Get base JSON length
            let json_len = serde_json::to_vec(&self.header)
                .map_err(|e| Error::InvalidFormat(format!("Failed to serialize header: {}", e)))?
                .len();

            // Add padding bytes to allow for future field expansion
            self.header_max_length = json_len + 64;

            // Write the padded header
            self.header
                .write(&mut self.writer, Some(self.header_max_length))?;

            // Reserve space for the frame index table offset
            self.writer.write_u64::<LittleEndian>(0)?;

            self.header_written = true;
        }
        Ok(())
    }

    /// Write a single frame
    pub fn write_frame(&mut self, frame: Frame) -> Result<()> {
        self.ensure_header_written()?;

        // Ensure the first frame is always a full frame
        let is_full_frame = self.num_frames == 0
            || (self.header.full_frame_interval > 0
                && self
                    .num_frames
                    .is_multiple_of(self.header.full_frame_interval as u64));

        // Store the current position for the frame offset table
        let frame_pos = self.writer.stream_position()?;
        self.frame_offsets.push(frame_pos);

        // Process and write the frame
        if !is_full_frame && self.reference_frame.is_some() {
            // Use delta encoding for positions, velocities, and forces
            let delta = frame
                .data
                .delta_from(self.reference_frame.as_ref().unwrap())?;

            // Create a delta frame with the original metadata
            let delta_frame = Frame {
                data: delta,
                is_full_frame: false,
                ..frame.clone()
            };

            // Add a flag to indicate this is a delta frame
            self.writer.write_u8(0u8)?; // 0 = delta frame

            // Create temporary compression settings with half precision for delta frames
            let mut delta_settings = self.compression_settings.clone();
            delta_settings.precision = PrecisionMode::Half;

            // Write the delta frame
            delta_frame.write(
                &mut self.writer,
                &delta_settings,
                &self.header.extra_datasets,
            )?;

            // Only update reference frame for delta frames
            self.reference_frame = Some(frame.data);
        } else {
            // Write a full frame (either first frame or at specified interval)
            // Add a flag to indicate this is a full frame
            self.writer.write_u8(1u8)?; // 1 = full frame

            // Create temporary compression settings with full precision for full frames
            let mut full_settings = self.compression_settings.clone();
            full_settings.precision = PrecisionMode::Full;

            // Write the full frame
            frame.write(
                &mut self.writer,
                &full_settings,
                &self.header.extra_datasets,
            )?;

            // Only update reference frame for full frames
            self.reference_frame = Some(frame.data);
        }

        // Count the frame
        self.num_frames += 1;

        Ok(())
    }

    /// Write multiple frames sequentially
    pub fn write_frames(&mut self, frames: Vec<Frame>) -> Result<()> {
        // Process frames sequentially to maintain proper reference frame state
        for frame in frames {
            self.write_frame(frame)?;
        }
        Ok(())
    }

    /// Finish writing the file, write the index and update the header
    pub fn finish(mut self) -> Result<()> {
        self.ensure_header_written()?;

        // Remember the position of the frame index table
        let index_pos = self.writer.stream_position()?;

        // Write the frame index table
        self.writer.write_u64::<LittleEndian>(self.num_frames)?;
        for offset in &self.frame_offsets {
            self.writer.write_u64::<LittleEndian>(*offset)?;
        }

        // Update the header information
        self.header.num_frames = self.num_frames;
        self.header.frame_index_offset = index_pos;

        // Go back to the beginning of the file and rewrite the header with identical length
        self.writer.seek(SeekFrom::Start(0))?;
        self.header
            .write(&mut self.writer, Some(self.header_max_length))?;

        // Write the frame index table offset
        self.writer.write_u64::<LittleEndian>(index_pos)?;

        // Flush the writer
        self.writer.flush()?;

        Ok(())
    }
}

/// Reader for TACO trajectory files
pub struct Reader {
    /// The file reader
    reader: BufReader<File>,

    /// Header information
    header: Header,

    /// Original length of the header JSON (used for appending)
    header_length: usize,

    /// Frame index table
    frame_offsets: Vec<u64>,

    /// Reference frame for delta decoding
    reference_frame: Option<FrameData>,

    /// Current frame number
    current_frame: u64,

    /// File path (stored for creating additional readers)
    file_path: std::path::PathBuf,
}

impl Reader {
    /// Open a TACO file for reading
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
        // Store the path
        let file_path = path.as_ref().to_path_buf();

        // Open the file
        let file = File::open(&file_path)?;
        // Use a larger buffer (1MB)
        let mut reader = BufReader::with_capacity(1024 * 1024, file);

        // Read the header
        let (header, header_length) = Header::read(&mut reader)?;

        for constraint in &header.constraints {
            if let Constraint::Unknown { name: _, .. } = constraint {
                // Warning: Unknown constraint found in header
            }
        }

        // Read the frame index table offset
        let frame_index_offset = reader.read_u64::<LittleEndian>()?;

        // Jump to the frame index table
        reader.seek(SeekFrom::Start(frame_index_offset))?;

        // Read the number of frames
        let num_frames = reader.read_u64::<LittleEndian>()?;

        if num_frames != header.num_frames {
            return Err(Error::InvalidFormat(format!(
                "Header says {} frames but index has {} frames",
                header.num_frames, num_frames
            )));
        }

        // Read the frame offsets
        let mut frame_offsets = Vec::with_capacity(num_frames as usize);
        for _ in 0..num_frames {
            frame_offsets.push(reader.read_u64::<LittleEndian>()?);
        }

        Ok(Reader {
            reader,
            header,
            header_length,
            frame_offsets,
            reference_frame: None,
            current_frame: u64::MAX, // No frame has been read yet
            file_path,
        })
    }

    /// Get the header information
    pub fn header(&self) -> &Header {
        &self.header
    }

    /// Get the number of frames in the trajectory
    pub fn num_frames(&self) -> u64 {
        self.header.num_frames
    }

    /// Get the number of atoms in the trajectory
    pub fn num_atoms(&self) -> u32 {
        self.header.num_atoms
    }

    /// Read a specific frame by index
    pub fn read_frame(&mut self, frame_index: u64) -> Result<Frame> {
        if frame_index >= self.frame_offsets.len() as u64 {
            return Err(Error::FrameAccess(format!(
                "Frame index {} out of bounds (max: {})",
                frame_index,
                self.frame_offsets.len() - 1
            )));
        }

        // Determine if we need to read keyframes leading up to this frame
        let full_frame_interval = self.header.full_frame_interval as u64;
        let is_keyframe =
            full_frame_interval == 0 || frame_index.is_multiple_of(full_frame_interval);

        // If this isn't a keyframe, we need to ensure we have the correct reference frame
        // by reading frames sequentially from the keyframe
        if !is_keyframe {
            // Find the preceding keyframe
            let keyframe_index = (frame_index / full_frame_interval) * full_frame_interval;

            // Check if we need to rebuild the reference frame chain
            let need_rebuild = self.reference_frame.is_none()
                || self.current_frame != frame_index - 1
                || (self.current_frame / full_frame_interval)
                    != (frame_index / full_frame_interval);

            if need_rebuild {
                // Read the keyframe first to establish the reference
                let _keyframe = self.read_frame(keyframe_index)?;

                // Now read all frames sequentially from keyframe+1 to frame_index-1
                for idx in (keyframe_index + 1)..frame_index {
                    let _intermediate = self.read_frame(idx)?;
                }
            }
        }

        // Seek to the frame position
        self.reader
            .seek(SeekFrom::Start(self.frame_offsets[frame_index as usize]))?;

        // Read frame flag to determine if this is a full or delta frame
        let frame_type_flag = self.reader.read_u8()?;
        let is_full_frame = frame_type_flag != 0;

        // First frame should always be a full frame
        if self.header.full_frame_interval > 0
            && frame_index.is_multiple_of(self.header.full_frame_interval as u64)
            && !is_full_frame
        {
            // The first frame should be a full frame
            let mut frame_settings = self.header.compression_settings.clone();
            frame_settings.precision = PrecisionMode::Full;

            // Read the frame data with full precision
            let mut frame = Frame::read(
                &mut self.reader,
                &frame_settings,
                &self.header.extra_datasets,
            )?;
            frame.is_full_frame = true;
            frame.frame_number = frame_index;

            // Store as reference frame
            self.reference_frame = Some(frame.data.clone());
            self.current_frame = frame_index;

            return Ok(frame);
        }

        // Create appropriate compression settings based on frame type
        let mut frame_settings = self.header.compression_settings.clone();
        if is_full_frame {
            frame_settings.precision = PrecisionMode::Full;
        } else {
            // For delta frames, use the default precision modes (typically half)
            frame_settings.precision = PrecisionMode::Half;
        }

        // Read the frame with the appropriate settings
        let mut frame = match Frame::read(
            &mut self.reader,
            &frame_settings,
            &self.header.extra_datasets,
        ) {
            Ok(f) => f,
            Err(e) => {
                return Err(Error::FrameAccess(format!(
                    "Failed to read frame {frame_index}: {e}"
                )));
            }
        };

        // Set the frame number and full frame flag
        frame.frame_number = frame_index;
        frame.is_full_frame = is_full_frame;

        // If it's a delta frame, apply delta to the reference frame
        if !is_full_frame {
            if let Some(reference) = &self.reference_frame {
                let absolute_data = reference.apply_delta(&frame.data)?;
                frame.data = absolute_data;
            } else {
                return Err(Error::FrameAccess(format!(
                    "Missing reference frame for delta frame {frame_index}"
                )));
            }
        }

        // Only clone if next frame will need it
        self.reference_frame = Some(frame.data.clone());
        self.current_frame = frame_index;

        Ok(frame)
    }

    /// Read the next frame sequentially
    pub fn next_frame(&mut self) -> Result<Option<Frame>> {
        if self.current_frame == u64::MAX {
            // First read, start from frame 0
            if self.num_frames() == 0 {
                return Ok(None);
            }
            let frame = self.read_frame(0)?;
            return Ok(Some(frame));
        }

        if self.current_frame >= self.num_frames() - 1 {
            return Ok(None);
        }

        let frame = self.read_frame(self.current_frame + 1)?;
        Ok(Some(frame))
    }

    /// Reset the reader to the beginning
    pub fn reset(&mut self) -> Result<()> {
        self.current_frame = u64::MAX;
        self.reference_frame = None;
        Ok(())
    }

    /// Iterate over all frames
    pub fn iter_frames(&mut self) -> FrameIterator<'_> {
        self.reset().unwrap();
        FrameIterator { reader: self }
    }

    /// Read a range of frames sequentially
    pub fn read_frame_range(&mut self, start: u64, end: u64) -> Result<Vec<Frame>> {
        let mut frames = Vec::with_capacity((end - start) as usize);
        for idx in start..end {
            let frame = self.read_frame(idx)?;
            frames.push(frame);
        }
        Ok(frames)
    }

    /// Read specific frames sequentially
    pub fn read_frames(&mut self, frame_indices: &[u64]) -> Result<Vec<Frame>> {
        if frame_indices.is_empty() {
            return Ok(Vec::new());
        }

        // Check all indices are valid
        for &idx in frame_indices {
            if idx >= self.frame_offsets.len() as u64 {
                return Err(Error::FrameAccess(format!(
                    "Frame index {} out of bounds (max: {})",
                    idx,
                    self.frame_offsets.len() - 1
                )));
            }
        }

        // Sort indices to read sequentially for proper delta handling
        let mut sorted_indices = frame_indices.to_vec();
        sorted_indices.sort_unstable();

        // Read frames in order
        let mut frames = Vec::with_capacity(frame_indices.len());

        // First make sure we have the right reference frame
        let first_idx = sorted_indices[0];
        let full_frame_interval = self.header.full_frame_interval as u64;

        // If first frame is not a keyframe, we need to read from the previous keyframe
        if full_frame_interval > 0
            && !first_idx.is_multiple_of(full_frame_interval)
            && first_idx > 0
        {
            let keyframe_idx = (first_idx / full_frame_interval) * full_frame_interval;

            // Read the keyframe first to establish reference
            self.read_frame(keyframe_idx)?;
        }

        // Now read the requested frames in sequence
        let mut result_frames = Vec::with_capacity(frame_indices.len());

        // Read in the sorted order to maintain delta references
        for &idx in &sorted_indices {
            let frame = self.read_frame(idx)?;
            frames.push(frame);
        }

        // Reorder frames to match the requested order
        for &idx in frame_indices {
            // Find this frame in our sorted results
            if let Some(position) = sorted_indices.iter().position(|&x| x == idx) {
                result_frames.push(frames[position].clone());
            }
        }

        Ok(result_frames)
    }

    /// Get the file path of the open file
    pub fn path(&self) -> &std::path::Path {
        &self.file_path
    }
}

/// Iterator over frames in a TACO trajectory
pub struct FrameIterator<'a> {
    reader: &'a mut Reader,
}

impl Iterator for FrameIterator<'_> {
    type Item = Result<Frame>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.reader.next_frame() {
            Ok(Some(frame)) => Some(Ok(frame)),
            Ok(None) => None,
            Err(e) => Some(Err(e)),
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        if self.reader.current_frame == u64::MAX {
            let n = self.reader.num_frames();
            return (n as usize, Some(n as usize));
        }
        let remaining = self
            .reader
            .num_frames()
            .saturating_sub(self.reader.current_frame + 1);
        (remaining as usize, Some(remaining as usize))
    }
}

#[cfg(feature = "async")]
#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        AtomMetadata, CompressionSettings, Constraint, Frame, FrameData, SimulationMetadata,
    };
    use ndarray::Array2;
    use tempfile::tempdir;

    #[test]
    fn test_append_basic() {
        let dir = tempdir().expect("Failed to create temp dir");
        let file_path = dir.path().join("test_append.taco");

        // Create initial file with 5 frames
        let sim_metadata = SimulationMetadata::default();
        let atom_metadata = AtomMetadata::default();
        let compression_settings = CompressionSettings::default();

        let mut writer = Writer::create(
            &file_path,
            3,
            0.001,
            sim_metadata.clone(),
            atom_metadata.clone(),
            compression_settings.clone(),
        )
        .expect("Failed to create writer");

        // Write initial 5 frames
        let mut positions = Array2::<f32>::zeros((3, 3));
        for i in 0..5 {
            for j in 0..3 {
                for k in 0..3 {
                    positions[[j, k]] = (i as f32) * 0.1 + (j as f32) + (k as f32) * 0.01;
                }
            }
            let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
            writer.write_frame(frame).expect("Failed to write frame");
        }
        writer.finish().expect("Failed to finish");

        // Now append 3 more frames
        let mut writer = Writer::append(&file_path).expect("Failed to open for append");

        for i in 5..8 {
            for j in 0..3 {
                for k in 0..3 {
                    positions[[j, k]] = (i as f32) * 0.1 + (j as f32) + (k as f32) * 0.01;
                }
            }
            let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
            writer.write_frame(frame).expect("Failed to write frame");
        }
        writer.finish().expect("Failed to finish");

        // Verify we can read all 8 frames
        let mut reader = Reader::open(&file_path).expect("Failed to open reader");
        assert_eq!(reader.num_frames(), 8);

        for i in 0..8 {
            let frame = reader.read_frame(i).expect("Failed to read frame");
            assert_eq!(frame.frame_number, i);

            if let Some(pos) = &frame.data.positions {
                for j in 0..3 {
                    for k in 0..3 {
                        let expected = (i as f32) * 0.1 + (j as f32) + (k as f32) * 0.01;
                        let actual = pos[[j, k]];
                        assert!(
                            (actual - expected).abs() < 0.001,
                            "Frame {}, atom {}, coord {}: expected {}, got {}",
                            i,
                            j,
                            k,
                            expected,
                            actual
                        );
                    }
                }
            }
        }
    }

    #[test]
    fn test_append_multiple_times() {
        let dir = tempdir().expect("Failed to create temp dir");
        let file_path = dir.path().join("test_append_multiple.taco");

        let sim_metadata = SimulationMetadata::default();
        let atom_metadata = AtomMetadata::default();
        let compression_settings = CompressionSettings::default();

        // Create initial file with 2 frames
        let mut writer = Writer::create(
            &file_path,
            2,
            0.001,
            sim_metadata,
            atom_metadata,
            compression_settings,
        )
        .expect("Failed to create writer");

        let mut positions = Array2::<f32>::zeros((2, 3));
        for i in 0..2 {
            for j in 0..2 {
                for k in 0..3 {
                    positions[[j, k]] = (i as f32) * 10.0 + (j as f32) + (k as f32) * 0.1;
                }
            }
            let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
            writer.write_frame(frame).expect("Failed to write frame");
        }
        writer.finish().expect("Failed to finish");

        // Append 2 more frames
        let mut writer = Writer::append(&file_path).expect("Failed to append");
        for i in 2..4 {
            for j in 0..2 {
                for k in 0..3 {
                    positions[[j, k]] = (i as f32) * 10.0 + (j as f32) + (k as f32) * 0.1;
                }
            }
            let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
            writer.write_frame(frame).expect("Failed to write frame");
        }
        writer.finish().expect("Failed to finish");

        // Append 3 more frames
        let mut writer = Writer::append(&file_path).expect("Failed to append");
        for i in 4..7 {
            for j in 0..2 {
                for k in 0..3 {
                    positions[[j, k]] = (i as f32) * 10.0 + (j as f32) + (k as f32) * 0.1;
                }
            }
            let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
            writer.write_frame(frame).expect("Failed to write frame");
        }
        writer.finish().expect("Failed to finish");

        // Verify all 7 frames
        let mut reader = Reader::open(&file_path).expect("Failed to open reader");
        assert_eq!(reader.num_frames(), 7);

        for i in 0..7 {
            let frame = reader.read_frame(i).expect("Failed to read frame");
            assert_eq!(frame.frame_number, i);

            if let Some(pos) = &frame.data.positions {
                for j in 0..2 {
                    for k in 0..3 {
                        let expected = (i as f32) * 10.0 + (j as f32) + (k as f32) * 0.1;
                        let actual = pos[[j, k]];
                        assert!(
                            (actual - expected).abs() < 0.01,
                            "Frame {}, atom {}, coord {}: expected {}, got {}",
                            i,
                            j,
                            k,
                            expected,
                            actual
                        );
                    }
                }
            }
        }
    }

    #[test]
    fn test_append_with_full_frame_interval() {
        let dir = tempdir().expect("Failed to create temp dir");
        let file_path = dir.path().join("test_append_full_interval.taco");

        let sim_metadata = SimulationMetadata::default();
        let atom_metadata = AtomMetadata::default();
        let compression_settings = CompressionSettings::default();

        // Create header with full_frame_interval = 3
        let mut header = Header::new(2, 0.001, sim_metadata, atom_metadata, compression_settings);
        header.full_frame_interval = 3;

        let mut writer =
            Writer::create_with_header(&file_path, header).expect("Failed to create writer");

        // Write initial 4 frames (0, 1, 2, 3 - where 0 and 3 are full frames)
        let mut positions = Array2::<f32>::zeros((2, 3));
        for i in 0..4 {
            for j in 0..2 {
                for k in 0..3 {
                    positions[[j, k]] = (i as f32) + (j as f32) * 0.1 + (k as f32) * 0.01;
                }
            }
            let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
            writer.write_frame(frame).expect("Failed to write frame");
        }
        writer.finish().expect("Failed to finish");

        // Append 5 more frames (4, 5, 6, 7, 8 - where 6 is a full frame)
        let mut writer = Writer::append(&file_path).expect("Failed to append");
        for i in 4..9 {
            for j in 0..2 {
                for k in 0..3 {
                    positions[[j, k]] = (i as f32) + (j as f32) * 0.1 + (k as f32) * 0.01;
                }
            }
            let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
            writer.write_frame(frame).expect("Failed to write frame");
        }
        writer.finish().expect("Failed to finish");

        // Verify all frames
        let mut reader = Reader::open(&file_path).expect("Failed to open reader");
        assert_eq!(reader.num_frames(), 9);
        assert_eq!(reader.header().full_frame_interval, 3);

        for i in 0..9 {
            let frame = reader.read_frame(i).expect("Failed to read frame");
            assert_eq!(frame.frame_number, i);

            // Check if full frame at correct intervals
            let expected_full = i % 3 == 0;
            assert_eq!(
                frame.is_full_frame,
                expected_full,
                "Frame {} should be {} but was {}",
                i,
                if expected_full { "full" } else { "delta" },
                if frame.is_full_frame { "full" } else { "delta" }
            );

            if let Some(pos) = &frame.data.positions {
                for j in 0..2 {
                    for k in 0..3 {
                        let expected = (i as f32) + (j as f32) * 0.1 + (k as f32) * 0.01;
                        let actual = pos[[j, k]];
                        assert!(
                            (actual - expected).abs() < 0.001,
                            "Frame {}, atom {}, coord {}: expected {}, got {}",
                            i,
                            j,
                            k,
                            expected,
                            actual
                        );
                    }
                }
            }
        }
    }

    #[test]
    fn test_append_empty_file() {
        let dir = tempdir().expect("Failed to create temp dir");
        let file_path = dir.path().join("test_append_empty.taco");

        let sim_metadata = SimulationMetadata::default();
        let atom_metadata = AtomMetadata::default();
        let compression_settings = CompressionSettings::default();

        // Create empty file (0 frames)
        let writer = Writer::create(
            &file_path,
            2,
            0.001,
            sim_metadata,
            atom_metadata,
            compression_settings,
        )
        .expect("Failed to create writer");
        writer.finish().expect("Failed to finish");

        // Append frames to empty file
        let mut writer = Writer::append(&file_path).expect("Failed to append");
        let mut positions = Array2::<f32>::zeros((2, 3));
        for i in 0..3 {
            for j in 0..2 {
                for k in 0..3 {
                    positions[[j, k]] = (i as f32) + (j as f32) * 0.1;
                }
            }
            let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
            writer.write_frame(frame).expect("Failed to write frame");
        }
        writer.finish().expect("Failed to finish");

        // Verify
        let mut reader = Reader::open(&file_path).expect("Failed to open reader");
        assert_eq!(reader.num_frames(), 3);

        for i in 0..3 {
            let frame = reader.read_frame(i).expect("Failed to read frame");
            assert_eq!(frame.frame_number, i);
        }
    }

    #[test]
    fn test_write_read_constraints() {
        let dir = tempdir().expect("Failed to create temp dir");
        let file_path = dir.path().join("test_constraints_roundtrip.taco");

        let sim_metadata = SimulationMetadata::default();
        let atom_metadata = AtomMetadata::default();
        let compression_settings = CompressionSettings::default();

        let mut writer = Writer::create(
            &file_path,
            3,
            0.001,
            sim_metadata,
            atom_metadata,
            compression_settings,
        )
        .expect("Failed to create writer");

        writer.set_constraints(vec![Constraint::FixAtoms {
            indices: vec![0, 2],
        }]);

        let positions = Array2::<f32>::zeros((3, 3));
        let frame = Frame::new(0, 0.0, FrameData::new(positions));
        writer.write_frame(frame).expect("Failed to write frame");
        writer.finish().expect("Failed to finish");

        let reader = Reader::open(&file_path).expect("Failed to open reader");
        let header = reader.header();

        assert_eq!(header.constraints.len(), 1);
        match &header.constraints[0] {
            Constraint::FixAtoms { indices } => {
                assert_eq!(indices, &vec![0, 2]);
            }
            _ => panic!("Expected FixAtoms constraint"),
        }
    }
}

#[cfg(feature = "async")]
pub mod r#async {
    //! Async I/O operations for TACO format
    //!
    //! This module provides async versions of the TACO readers and writers
    //! for non-blocking I/O operations. This is only available when the
    //! "async" feature is enabled.

    use super::*;

    /// Async wrapper around the sync Writer
    pub struct AsyncWriter {
        inner: Option<Writer>,
    }

    impl AsyncWriter {
        /// Create a new async TACO writer
        pub async fn create<P: AsRef<std::path::Path>>(
            path: P,
            num_atoms: u32,
            time_step: f64,
            simulation_metadata: SimulationMetadata,
            atom_metadata: AtomMetadata,
            compression_settings: CompressionSettings,
        ) -> Result<Self> {
            let path_buf = path.as_ref().to_path_buf();

            // Create the writer in a blocking task to avoid blocking the async runtime
            let writer = tokio::task::spawn_blocking({
                let path_buf = path_buf.clone();
                move || {
                    Writer::create(
                        path_buf,
                        num_atoms,
                        time_step,
                        simulation_metadata,
                        atom_metadata,
                        compression_settings,
                    )
                }
            })
            .await
            .map_err(|e| Error::Io(std::io::Error::other(e)))?;

            Ok(Self {
                inner: Some(writer?),
            })
        }

        /// Set constraints for the header (call before first write_frame)
        pub fn set_constraints(&mut self, constraints: Vec<Constraint>) {
            if let Some(writer) = self.inner.as_mut() {
                writer.set_constraints(constraints);
            } else {
                eprintln!("Warning: AsyncWriter already finished; constraints not applied");
            }
        }

        /// Set periodic boundary conditions for the header (call before first write_frame)
        pub fn set_pbc(&mut self, pbc: [bool; 3]) {
            if let Some(writer) = self.inner.as_mut() {
                writer.set_pbc(pbc);
            } else {
                eprintln!("Warning: AsyncWriter already finished; PBC not applied");
            }
        }

        /// Open an existing TACO file for appending frames (async)
        pub async fn append<P: AsRef<std::path::Path>>(path: P) -> Result<Self> {
            let path_buf = path.as_ref().to_path_buf();

            let writer = tokio::task::spawn_blocking(move || Writer::append(path_buf))
                .await
                .map_err(|e| Error::Io(std::io::Error::other(e)))?;

            Ok(Self {
                inner: Some(writer?),
            })
        }

        /// Write a frame to the trajectory (async)
        pub async fn write_frame(&mut self, frame: Frame) -> Result<()> {
            if let Some(mut writer) = self.inner.take() {
                let result = tokio::task::spawn_blocking(move || {
                    let result = writer.write_frame(frame);
                    (writer, result)
                })
                .await
                .map_err(|e| Error::Io(std::io::Error::other(e)))?;

                self.inner = Some(result.0);
                result.1
            } else {
                Err(Error::Io(std::io::Error::new(
                    std::io::ErrorKind::InvalidInput,
                    "Writer already finished",
                )))
            }
        }

        /// Write multiple frames to the trajectory (async)
        pub async fn write_frames(&mut self, frames: Vec<Frame>) -> Result<()> {
            if let Some(mut writer) = self.inner.take() {
                let result = tokio::task::spawn_blocking(move || {
                    let result = writer.write_frames(frames);
                    (writer, result)
                })
                .await
                .map_err(|e| Error::Io(std::io::Error::other(e)))?;

                self.inner = Some(result.0);
                result.1
            } else {
                Err(Error::Io(std::io::Error::new(
                    std::io::ErrorKind::InvalidInput,
                    "Writer already finished",
                )))
            }
        }

        /// Finish writing (async)
        pub async fn finish(mut self) -> Result<()> {
            if let Some(writer) = self.inner.take() {
                tokio::task::spawn_blocking(move || writer.finish())
                    .await
                    .map_err(|e| Error::Io(std::io::Error::other(e)))?
            } else {
                Err(Error::Io(std::io::Error::new(
                    std::io::ErrorKind::InvalidInput,
                    "Writer already finished",
                )))
            }
        }
    }

    /// Async wrapper around the sync Reader
    pub struct AsyncReader {
        inner: Option<Reader>,
    }

    impl AsyncReader {
        /// Open a TACO file for async reading
        pub async fn open<P: AsRef<std::path::Path>>(path: P) -> Result<Self> {
            let path_buf = path.as_ref().to_path_buf();

            let reader = tokio::task::spawn_blocking(move || Reader::open(path_buf))
                .await
                .map_err(|e| Error::Io(std::io::Error::other(e)))?;

            Ok(Self {
                inner: Some(reader?),
            })
        }

        /// Read a specific frame by index (async)
        pub async fn read_frame(&mut self, frame_index: u64) -> Result<Frame> {
            if let Some(mut reader) = self.inner.take() {
                let result = tokio::task::spawn_blocking(move || {
                    let result = reader.read_frame(frame_index);
                    (reader, result)
                })
                .await
                .map_err(|e| Error::Io(std::io::Error::other(e)))?;

                self.inner = Some(result.0);
                result.1
            } else {
                Err(Error::Io(std::io::Error::new(
                    std::io::ErrorKind::InvalidInput,
                    "Reader has been consumed",
                )))
            }
        }

        /// Read the next frame sequentially (async)
        pub async fn next_frame(&mut self) -> Result<Option<Frame>> {
            if let Some(mut reader) = self.inner.take() {
                let result = tokio::task::spawn_blocking(move || {
                    let result = reader.next_frame();
                    (reader, result)
                })
                .await
                .map_err(|e| Error::Io(std::io::Error::other(e)))?;

                self.inner = Some(result.0);
                result.1
            } else {
                Err(Error::Io(std::io::Error::new(
                    std::io::ErrorKind::InvalidInput,
                    "Reader has been consumed",
                )))
            }
        }

        /// Reset the reader to the beginning (async)
        pub async fn reset(&mut self) -> Result<()> {
            if let Some(mut reader) = self.inner.take() {
                let result = tokio::task::spawn_blocking(move || {
                    let result = reader.reset();
                    (reader, result)
                })
                .await
                .map_err(|e| Error::Io(std::io::Error::other(e)))?;

                self.inner = Some(result.0);
                result.1
            } else {
                Err(Error::Io(std::io::Error::new(
                    std::io::ErrorKind::InvalidInput,
                    "Reader has been consumed",
                )))
            }
        }

        /// Read a range of frames (async)
        pub async fn read_frame_range(&mut self, start: u64, end: u64) -> Result<Vec<Frame>> {
            if let Some(mut reader) = self.inner.take() {
                let result = tokio::task::spawn_blocking(move || {
                    let result = reader.read_frame_range(start, end);
                    (reader, result)
                })
                .await
                .map_err(|e| Error::Io(std::io::Error::other(e)))?;

                self.inner = Some(result.0);
                result.1
            } else {
                Err(Error::Io(std::io::Error::new(
                    std::io::ErrorKind::InvalidInput,
                    "Reader has been consumed",
                )))
            }
        }

        /// Read specific frames (async)
        pub async fn read_frames(&mut self, frame_indices: &[u64]) -> Result<Vec<Frame>> {
            if let Some(mut reader) = self.inner.take() {
                let frame_indices = frame_indices.to_vec();
                let result = tokio::task::spawn_blocking(move || {
                    let result = reader.read_frames(&frame_indices);
                    (reader, result)
                })
                .await
                .map_err(|e| Error::Io(std::io::Error::other(e)))?;

                self.inner = Some(result.0);
                result.1
            } else {
                Err(Error::Io(std::io::Error::new(
                    std::io::ErrorKind::InvalidInput,
                    "Reader has been consumed",
                )))
            }
        }

        /// Get the header information
        pub fn header(&self) -> Option<&Header> {
            self.inner.as_ref().map(|r| r.header())
        }

        /// Get the number of frames
        pub fn num_frames(&self) -> u64 {
            self.inner.as_ref().map(|r| r.num_frames()).unwrap_or(0)
        }

        /// Get the number of atoms
        pub fn num_atoms(&self) -> u32 {
            self.inner.as_ref().map(|r| r.num_atoms()).unwrap_or(0)
        }

        /// Get the file path
        pub fn path(&self) -> Option<&std::path::Path> {
            self.inner.as_ref().map(|r| r.path())
        }
    }
}
