json_archive/

archive_context.rs

// json-archive is a tool for tracking JSON file changes over time
// Copyright (C) 2025  Peoples Grocers LLC
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
// To purchase a license under different terms contact admin@peoplesgrocers.com
// To request changes, report bugs, or give user feedback contact
// marxism@peoplesgrocers.com
//

//! Archive write context and shared observation writing logic.
//!
//! This module provides:
//! - `WriteContext`: A struct that holds the state needed to write observations
//! - `write_observations`: The shared logic for diffing JSON files and writing events
//!
//! The key insight is that both create and append operations share the same
//! core logic once they've set up their initial state and writer.

use chrono::{DateTime, Utc};
use serde_json::Value;
use std::io::Write;
use std::path::{Path, PathBuf};
use uuid::Uuid;

use crate::atomic_file::atomic_replace_file;
use crate::detection::CompressionFormat;
use crate::diagnostics::{Diagnostic, DiagnosticCode, DiagnosticCollector};
use crate::diff;
use crate::events::{Event, Observation};

/// Strategy for finishing the write operation.
#[derive(Debug, Clone)]
pub enum FinishStrategy {
    /// Just flush the writer. Used for:
    /// - Creating new archives
    /// - Appending to uncompressed archives (same file)
    FlushOnly,

    /// Atomic replace: swap temp file with original. Used for:
    /// - Appending to compressed archives (rewrite strategy)
    AtomicReplace {
        temp_path: PathBuf,
        output_path: PathBuf,
    },
}

/// Context for writing observations to an archive.
///
/// This struct is the result of the "setup phase" for both create and append
/// operations. Once you have a WriteContext, you can use `write_observations`
/// to add new states, then call `finish` to complete the operation.
pub struct WriteContext<W: Write> {
    /// The writer to output JSON lines to.
    pub writer: W,

    /// Current state of the archive (used for diffing).
    pub current_state: Value,

    /// Number of observations already in the archive.
    pub observation_count: usize,

    /// Optional interval for writing snapshots.
    pub snapshot_interval: Option<usize>,

    /// How to finish the write operation.
    pub finish_strategy: FinishStrategy,

    /// Diagnostics collected during setup (e.g., warnings from reading existing archive).
    pub diagnostics: DiagnosticCollector,
}

impl<W: Write> WriteContext<W> {
    /// Create a new write context.
    pub fn new(
        writer: W,
        current_state: Value,
        observation_count: usize,
        snapshot_interval: Option<usize>,
        finish_strategy: FinishStrategy,
    ) -> Self {
        Self {
            writer,
            current_state,
            observation_count,
            snapshot_interval,
            finish_strategy,
            diagnostics: DiagnosticCollector::new(),
        }
    }

    /// Create a write context with existing diagnostics.
    pub fn with_diagnostics(
        writer: W,
        current_state: Value,
        observation_count: usize,
        snapshot_interval: Option<usize>,
        finish_strategy: FinishStrategy,
        diagnostics: DiagnosticCollector,
    ) -> Self {
        Self {
            writer,
            current_state,
            observation_count,
            snapshot_interval,
            finish_strategy,
            diagnostics,
        }
    }

    /// Write observations for a list of JSON files.
    ///
    /// For each file:
    /// 1. Reads and parses the JSON
    /// 2. Diffs against current state
    /// 3. Writes observation events
    /// 4. Optionally writes a snapshot if interval is reached
    /// 5. Updates current state
    ///
    /// Returns the number of observations written.
    pub fn write_observations<P: AsRef<Path>>(
        &mut self,
        files: &[P],
    ) -> Result<usize, Vec<Diagnostic>> {
        let mut observations_written = 0;

        for file_path in files.iter() {
            let file_path = file_path.as_ref();

            // Write comment marking which file we're processing
            if let Err(e) = writeln!(self.writer, "# Processing file: {}", file_path.display()) {
                return Err(vec![Diagnostic::fatal(
                    DiagnosticCode::PathNotFound,
                    format!("I couldn't write to the output: {}", e),
                )]);
            }

            // Get file modification time for the observation timestamp
            let file_mtime = get_file_mtime(file_path)?;

            // Read and parse new state
            let content = std::fs::read_to_string(file_path).map_err(|e| {
                vec![Diagnostic::fatal(
                    DiagnosticCode::PathNotFound,
                    format!("I couldn't read the input file '{}': {}", file_path.display(), e),
                )]
            })?;

            let new_state: Value = serde_json::from_str(&content).map_err(|e| {
                vec![Diagnostic::fatal(
                    DiagnosticCode::InvalidEventJson,
                    format!("I couldn't parse '{}' as JSON: {}", file_path.display(), e),
                )
                .with_advice("Make sure the file contains valid JSON.".to_string())]
            })?;

            // Generate diff and create observation
            let observation_id = format!("obs-{}", Uuid::new_v4());
            let diff_events = diff::diff(&self.current_state, &new_state, "", &observation_id);

            // Skip if no changes
            if diff_events.is_empty() {
                continue;
            }

            // Create and write observation
            let mut observation = Observation::new(observation_id, file_mtime);
            for event in diff_events {
                observation.add_event(event);
            }

            self.write_observation(observation)?;
            observations_written += 1;
            self.observation_count += 1;

            // Check if we should write a snapshot
            if self.should_write_snapshot() {
                self.write_snapshot(&new_state, file_mtime)?;
            }

            // Update current state for next iteration
            self.current_state = new_state;
        }

        Ok(observations_written)
    }

    /// Write a single observation's events to the output.
    fn write_observation(&mut self, observation: Observation) -> Result<(), Vec<Diagnostic>> {
        for event in observation.to_events() {
            let event_json = serde_json::to_string(&event).map_err(|e| {
                vec![Diagnostic::fatal(
                    DiagnosticCode::InvalidEventJson,
                    format!("I couldn't serialize an event to JSON: {}", e),
                )]
            })?;

            writeln!(self.writer, "{}", event_json).map_err(|e| {
                vec![Diagnostic::fatal(
                    DiagnosticCode::PathNotFound,
                    format!("I couldn't write to the output: {}", e),
                )]
            })?;
        }

        Ok(())
    }

    /// Check if we should write a snapshot based on observation count.
    fn should_write_snapshot(&self) -> bool {
        if let Some(interval) = self.snapshot_interval {
            self.observation_count > 0 && self.observation_count % interval == 0
        } else {
            false
        }
    }

    /// Write a snapshot event.
    fn write_snapshot(&mut self, state: &Value, timestamp: DateTime<Utc>) -> Result<(), Vec<Diagnostic>> {
        let snapshot_id = format!("snapshot-{}", Uuid::new_v4());
        let snapshot = Event::Snapshot {
            observation_id: snapshot_id,
            timestamp,
            object: state.clone(),
        };

        let snapshot_json = serde_json::to_string(&snapshot).map_err(|e| {
            vec![Diagnostic::fatal(
                DiagnosticCode::InvalidEventJson,
                format!("I couldn't serialize the snapshot to JSON: {}", e),
            )]
        })?;

        writeln!(self.writer, "{}", snapshot_json).map_err(|e| {
            vec![Diagnostic::fatal(
                DiagnosticCode::PathNotFound,
                format!("I couldn't write to the output: {}", e),
            )]
        })?;

        Ok(())
    }

    /// Finish the write operation.
    ///
    /// This flushes the writer and, for compressed append operations,
    /// performs the atomic file replacement.
    pub fn finish(mut self) -> Result<DiagnosticCollector, Vec<Diagnostic>> {
        // Flush the writer
        self.writer.flush().map_err(|e| {
            vec![Diagnostic::fatal(
                DiagnosticCode::PathNotFound,
                format!("I couldn't flush the output file: {}", e),
            )]
        })?;

        // Handle atomic replacement if needed
        match self.finish_strategy {
            FinishStrategy::FlushOnly => {
                // Nothing more to do
            }
            FinishStrategy::AtomicReplace { temp_path, output_path } => {
                atomic_replace_file(&output_path, &temp_path)?;
            }
        }

        Ok(self.diagnostics)
    }
}

/// Get the file modification time as a DateTime<Utc>.
fn get_file_mtime<P: AsRef<Path>>(path: P) -> Result<DateTime<Utc>, Vec<Diagnostic>> {
    let path = path.as_ref();
    let metadata = std::fs::metadata(path).map_err(|e| {
        vec![Diagnostic::fatal(
            DiagnosticCode::PathNotFound,
            format!("I couldn't get metadata for '{}': {}", path.display(), e),
        )]
    })?;

    let modified = metadata.modified().map_err(|e| {
        vec![Diagnostic::fatal(
            DiagnosticCode::PathNotFound,
            format!("I couldn't get modification time for '{}': {}", path.display(), e),
        )]
    })?;

    Ok(modified.into())
}

/// Encoder wrapper that provides a uniform interface for different compression formats.
///
/// This enum wraps the various compression encoders so we can treat them uniformly
/// in the append-to-compressed-archive flow.
#[cfg(feature = "compression")]
pub enum CompressedWriter {
    Gzip(flate2::write::GzEncoder<std::fs::File>),
    Zlib(flate2::write::ZlibEncoder<std::fs::File>),
    Zstd(zstd::stream::write::Encoder<'static, std::fs::File>),
    Brotli(brotli::CompressorWriter<std::fs::File>),
}

#[cfg(feature = "compression")]
impl Write for CompressedWriter {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        match self {
            CompressedWriter::Gzip(w) => w.write(buf),
            CompressedWriter::Zlib(w) => w.write(buf),
            CompressedWriter::Zstd(w) => w.write(buf),
            CompressedWriter::Brotli(w) => w.write(buf),
        }
    }

    fn flush(&mut self) -> std::io::Result<()> {
        match self {
            CompressedWriter::Gzip(w) => w.flush(),
            CompressedWriter::Zlib(w) => w.flush(),
            CompressedWriter::Zstd(w) => w.flush(),
            CompressedWriter::Brotli(w) => w.flush(),
        }
    }
}

#[cfg(feature = "compression")]
impl CompressedWriter {
    /// Create a new compressed writer for the given format and file.
    pub fn new(format: CompressionFormat, file: std::fs::File) -> Result<Self, Diagnostic> {
        use flate2::Compression;

        match format {
            CompressionFormat::Gzip => {
                Ok(CompressedWriter::Gzip(flate2::write::GzEncoder::new(file, Compression::default())))
            }
            CompressionFormat::Zlib => {
                Ok(CompressedWriter::Zlib(flate2::write::ZlibEncoder::new(file, Compression::default())))
            }
            CompressionFormat::Zstd => {
                let encoder = zstd::stream::write::Encoder::new(file, 0).map_err(|e| {
                    Diagnostic::fatal(
                        DiagnosticCode::PathNotFound,
                        format!("I couldn't create zstd encoder: {}", e),
                    )
                })?;
                Ok(CompressedWriter::Zstd(encoder))
            }
            CompressionFormat::Brotli => {
                Ok(CompressedWriter::Brotli(brotli::CompressorWriter::new(file, 4096, 11, 22)))
            }
            CompressionFormat::Deflate => {
                // Deflate is typically used within gzip/zlib, not standalone for files
                Err(Diagnostic::fatal(
                    DiagnosticCode::UnsupportedVersion,
                    "Standalone deflate compression is not supported for writing.".to_string(),
                ))
            }
            CompressionFormat::None => {
                Err(Diagnostic::fatal(
                    DiagnosticCode::UnsupportedVersion,
                    "CompressedWriter::new called with CompressionFormat::None".to_string(),
                ))
            }
        }
    }

    /// Finish compression and return any errors.
    ///
    /// This must be called before the file is closed to ensure all
    /// compressed data is flushed.
    pub fn finish(self) -> Result<(), Diagnostic> {
        match self {
            CompressedWriter::Gzip(w) => {
                w.finish().map_err(|e| {
                    Diagnostic::fatal(
                        DiagnosticCode::PathNotFound,
                        format!("I couldn't finish gzip compression: {}", e),
                    )
                })?;
            }
            CompressedWriter::Zlib(w) => {
                w.finish().map_err(|e| {
                    Diagnostic::fatal(
                        DiagnosticCode::PathNotFound,
                        format!("I couldn't finish zlib compression: {}", e),
                    )
                })?;
            }
            CompressedWriter::Zstd(w) => {
                w.finish().map_err(|e| {
                    Diagnostic::fatal(
                        DiagnosticCode::PathNotFound,
                        format!("I couldn't finish zstd compression: {}", e),
                    )
                })?;
            }
            CompressedWriter::Brotli(mut w) => {
                // Brotli doesn't have a finish() method, flush is sufficient
                w.flush().map_err(|e| {
                    Diagnostic::fatal(
                        DiagnosticCode::PathNotFound,
                        format!("I couldn't flush brotli compression: {}", e),
                    )
                })?;
            }
        }
        Ok(())
    }
}

/// A write context specifically for compressed output.
///
/// This wraps WriteContext to handle the finish() call properly for
/// compressed writers, which need to call finish() on the encoder
/// before the atomic file swap.
#[cfg(feature = "compression")]
pub struct CompressedWriteContext {
    /// The inner write context.
    inner: WriteContext<CompressedWriter>,
}

#[cfg(feature = "compression")]
impl CompressedWriteContext {
    /// Create a new compressed write context.
    pub fn new(
        writer: CompressedWriter,
        current_state: Value,
        observation_count: usize,
        snapshot_interval: Option<usize>,
        finish_strategy: FinishStrategy,
        diagnostics: DiagnosticCollector,
    ) -> Self {
        Self {
            inner: WriteContext::with_diagnostics(
                writer,
                current_state,
                observation_count,
                snapshot_interval,
                finish_strategy,
                diagnostics,
            ),
        }
    }

    /// Write observations for a list of JSON files.
    pub fn write_observations<P: AsRef<Path>>(
        &mut self,
        files: &[P],
    ) -> Result<usize, Vec<Diagnostic>> {
        self.inner.write_observations(files)
    }

    /// Write raw bytes to the output (used for copying existing archive content).
    pub fn write_raw(&mut self, bytes: &[u8]) -> Result<(), Vec<Diagnostic>> {
        self.inner.writer.write_all(bytes).map_err(|e| {
            vec![Diagnostic::fatal(
                DiagnosticCode::PathNotFound,
                format!("I couldn't write to the output: {}", e),
            )]
        })
    }

    /// Finish the write operation.
    ///
    /// This finishes the compression encoder, then performs any atomic
    /// file operations needed.
    pub fn finish(self) -> Result<DiagnosticCollector, Vec<Diagnostic>> {
        let finish_strategy = self.inner.finish_strategy.clone();
        let diagnostics = self.inner.diagnostics;

        // Finish compression first
        self.inner.writer.finish().map_err(|d| vec![d])?;

        // Then handle atomic replacement if needed
        match finish_strategy {
            FinishStrategy::FlushOnly => {
                // Nothing more to do
            }
            FinishStrategy::AtomicReplace { temp_path, output_path } => {
                atomic_replace_file(&output_path, &temp_path)?;
            }
        }

        Ok(diagnostics)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn test_write_context_single_observation() {
        let mut output = Vec::new();
        let initial_state = json!({"count": 0});

        {
            let mut ctx = WriteContext::new(
                &mut output,
                initial_state,
                0,
                None,
                FinishStrategy::FlushOnly,
            );

            // Create a temp file with new state
            let mut temp_file = tempfile::NamedTempFile::new().unwrap();
            std::io::Write::write_all(&mut temp_file, br#"{"count": 1}"#).unwrap();
            temp_file.flush().unwrap();

            let count = ctx.write_observations(&[temp_file.path()]).unwrap();
            assert_eq!(count, 1);
        }

        let output_str = String::from_utf8(output).unwrap();
        assert!(output_str.contains("# Processing file:"));
        assert!(output_str.contains("observe"));
        assert!(output_str.contains("change"));
        assert!(output_str.contains("/count"));
    }

    #[test]
    fn test_write_context_no_changes() {
        let mut output = Vec::new();
        let initial_state = json!({"count": 0});

        {
            let mut ctx = WriteContext::new(
                &mut output,
                initial_state,
                0,
                None,
                FinishStrategy::FlushOnly,
            );

            // Create a temp file with same state
            let mut temp_file = tempfile::NamedTempFile::new().unwrap();
            std::io::Write::write_all(&mut temp_file, br#"{"count": 0}"#).unwrap();
            temp_file.flush().unwrap();

            let count = ctx.write_observations(&[temp_file.path()]).unwrap();
            assert_eq!(count, 0);
        }

        let output_str = String::from_utf8(output).unwrap();
        // Should have comment but no events
        assert!(output_str.contains("# Processing file:"));
        assert!(!output_str.contains("observe"));
    }

    #[test]
    fn test_should_write_snapshot() {
        let output: Vec<u8> = Vec::new();

        // No interval set
        let ctx: WriteContext<Vec<u8>> = WriteContext::new(
            output.clone(),
            json!({}),
            5,
            None,
            FinishStrategy::FlushOnly,
        );
        assert!(!ctx.should_write_snapshot());

        // Interval of 2, at observation 4 (multiple of 2)
        let ctx: WriteContext<Vec<u8>> = WriteContext::new(
            output.clone(),
            json!({}),
            4,
            Some(2),
            FinishStrategy::FlushOnly,
        );
        assert!(ctx.should_write_snapshot());

        // Interval of 2, at observation 3 (not multiple of 2)
        let ctx: WriteContext<Vec<u8>> = WriteContext::new(
            output,
            json!({}),
            3,
            Some(2),
            FinishStrategy::FlushOnly,
        );
        assert!(!ctx.should_write_snapshot());
    }
}