zc2 0.0.13

//! Write-Ahead Log for crash recovery.
//!
//! Append-only JSON Lines file that tracks the lifecycle of each request:
//! Reserved → Executed → Committed/Failed
//!
//! On broker restart, uncommitted entries are replayed to ensure
//! no credits are lost if the broker crashes after worker execution.
//!
//! Optimizations:
//! - In-memory DashMap index for O(1) update_status lookups
//! - Batched fsync via WriteBuffer (flush every 100ms or 64 entries)

use chrono::{DateTime, Utc};
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use std::fs::{self, File, OpenOptions};
use std::io::{BufRead, BufReader, Write};
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Instant;

/// Max entries before flushing the write buffer to disk.
const WRITE_BUFFER_MAX_ENTRIES: usize = 64;
/// Max time (ms) before flushing the write buffer to disk.
const WRITE_BUFFER_MAX_AGE_MS: u128 = 100;

/// WAL entry status
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum WalStatus {
    Reserved,
    Executed,
    Committed,
    Failed,
}

/// A single WAL entry
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WalEntry {
    pub request_id: String,
    pub user_id: String,
    pub reservation_id: String,
    pub estimated_cost: f64,
    pub actual_cost: Option<f64>,
    pub worker_id: String,
    pub duration_ms: Option<f64>,
    pub timestamp: DateTime<Utc>,
    pub status: WalStatus,
}

/// Internal write buffer — collects serialized entries, flushes in batch.
struct WriteBuffer {
    lines: Vec<String>,
    last_flush: Instant,
}

impl WriteBuffer {
    fn new() -> Self {
        Self {
            lines: Vec::with_capacity(WRITE_BUFFER_MAX_ENTRIES),
            last_flush: Instant::now(),
        }
    }

    fn push(&mut self, line: String) {
        self.lines.push(line);
    }

    fn should_flush(&self) -> bool {
        self.lines.len() >= WRITE_BUFFER_MAX_ENTRIES
            || (!self.lines.is_empty()
                && self.last_flush.elapsed().as_millis() >= WRITE_BUFFER_MAX_AGE_MS)
    }

    fn drain(&mut self) -> Vec<String> {
        self.last_flush = Instant::now();
        std::mem::take(&mut self.lines)
    }

    fn is_empty(&self) -> bool {
        self.lines.is_empty()
    }
}

/// Write-Ahead Log
pub struct Wal {
    path: PathBuf,
    file: Mutex<File>,
    /// In-memory index: request_id → latest WalEntry (O(1) lookup).
    entries_index: DashMap<String, WalEntry>,
    /// Write buffer for batched fsync.
    write_buffer: Mutex<WriteBuffer>,
    /// Whether the index has been populated from disk.
    index_populated: AtomicBool,
}

/// Exposed for tests only — reads entries directly from disk (bypasses index).
#[cfg(test)]
pub fn read_all_from_disk(path: &std::path::Path) -> std::io::Result<Vec<WalEntry>> {
    let file = std::fs::File::open(path)?;
    let reader = std::io::BufReader::new(file);
    let mut entries = Vec::new();
    for line in reader.lines() {
        let line = line?;
        if !line.trim().is_empty() {
            if let Ok(e) = serde_json::from_str::<WalEntry>(&line) {
                entries.push(e);
            }
        }
    }
    Ok(entries)
}

impl Wal {
    /// Open or create a WAL file at the given path
    pub fn open<P: AsRef<Path>>(path: P) -> std::io::Result<Self> {
        let path = path.as_ref().to_path_buf();

        // Ensure parent directory exists
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)?;
        }

        let file = OpenOptions::new()
            .create(true)
            .append(true)
            .open(&path)?;

        let wal = Self {
            path,
            file: Mutex::new(file),
            entries_index: DashMap::new(),
            write_buffer: Mutex::new(WriteBuffer::new()),
            index_populated: AtomicBool::new(false),
        };

        // Populate index from existing WAL file
        wal.populate_index();

        Ok(wal)
    }

    /// Populate the in-memory index from the WAL file on disk.
    fn populate_index(&self) {
        if let Ok(entries) = self.read_all() {
            for entry in entries {
                self.entries_index.insert(entry.request_id.clone(), entry);
            }
        }
        self.index_populated.store(true, Ordering::Release);
    }

    /// Append a new entry to the WAL
    pub fn append(&self, entry: &WalEntry) -> std::io::Result<()> {
        let line = serde_json::to_string(entry).map_err(|e| {
            std::io::Error::new(std::io::ErrorKind::Other, e)
        })?;

        // Update in-memory index
        self.entries_index.insert(entry.request_id.clone(), entry.clone());

        // Add to write buffer
        let should_flush = {
            let mut buf = self.write_buffer.lock().map_err(|_| {
                std::io::Error::new(std::io::ErrorKind::Other, "Write buffer lock poisoned")
            })?;
            buf.push(line);
            buf.should_flush()
        };

        if should_flush {
            self.flush_buffer()?;
        }

        Ok(())
    }

    /// Update the status of an existing entry by appending a new line
    /// with the same request_id but updated status/fields.
    /// Uses the in-memory index for O(1) lookup instead of reading from disk.
    pub fn update_status(
        &self,
        request_id: &str,
        status: WalStatus,
        actual_cost: Option<f64>,
        duration_ms: Option<f64>,
    ) -> std::io::Result<()> {
        // Look up from in-memory index (O(1))
        let original = self.entries_index.get(request_id).map(|e| e.clone());

        if let Some(orig) = original {
            let updated = WalEntry {
                request_id: orig.request_id.clone(),
                user_id: orig.user_id.clone(),
                reservation_id: orig.reservation_id.clone(),
                estimated_cost: orig.estimated_cost,
                actual_cost: actual_cost.or(orig.actual_cost),
                worker_id: orig.worker_id.clone(),
                duration_ms: duration_ms.or(orig.duration_ms),
                timestamp: Utc::now(),
                status,
            };
            self.append(&updated)?;
        }
        Ok(())
    }

    /// Flush the write buffer to disk with a single fsync.
    pub fn flush_buffer(&self) -> std::io::Result<()> {
        let lines = {
            let mut buf = self.write_buffer.lock().map_err(|_| {
                std::io::Error::new(std::io::ErrorKind::Other, "Write buffer lock poisoned")
            })?;
            if buf.is_empty() {
                return Ok(());
            }
            buf.drain()
        };

        let mut file = self.file.lock().map_err(|_| {
            std::io::Error::new(std::io::ErrorKind::Other, "WAL lock poisoned")
        })?;

        for line in &lines {
            writeln!(file, "{}", line)?;
        }
        file.flush()?;
        file.sync_data()?;

        Ok(())
    }

    /// Read all entries from the WAL file
    fn read_all(&self) -> std::io::Result<Vec<WalEntry>> {
        let file = File::open(&self.path)?;
        let reader = BufReader::new(file);
        let mut entries = Vec::new();

        for line in reader.lines() {
            let line = line?;
            if line.trim().is_empty() {
                continue;
            }
            if let Ok(entry) = serde_json::from_str::<WalEntry>(&line) {
                entries.push(entry);
            }
        }
        Ok(entries)
    }

    /// Read uncommitted entries (the latest status for each request_id
    /// that is not Committed or Failed).
    /// Uses in-memory index when available.
    pub fn read_uncommitted(&self) -> std::io::Result<Vec<WalEntry>> {
        // Flush any buffered writes first
        let _ = self.flush_buffer();

        if self.index_populated.load(Ordering::Acquire) {
            // Fast path: read from DashMap index
            Ok(self.entries_index
                .iter()
                .filter(|e| e.status != WalStatus::Committed && e.status != WalStatus::Failed)
                .map(|e| e.value().clone())
                .collect())
        } else {
            // Fallback: read from disk
            let entries = self.read_all()?;
            let mut latest: std::collections::HashMap<String, WalEntry> =
                std::collections::HashMap::new();
            for entry in entries {
                latest.insert(entry.request_id.clone(), entry);
            }
            Ok(latest
                .into_values()
                .filter(|e| e.status != WalStatus::Committed && e.status != WalStatus::Failed)
                .collect())
        }
    }

    /// Path to the WAL file (for tests).
    #[cfg(test)]
    pub fn path(&self) -> &std::path::Path {
        &self.path
    }

    /// Compact the WAL: rewrite keeping only uncommitted entries.
    /// Also clears completed entries from the in-memory index.
    pub fn compact(&self) -> std::io::Result<()> {
        // Flush buffer before compacting
        self.flush_buffer()?;

        let uncommitted = self.read_uncommitted()?;

        // Write to temp file, then rename
        let tmp_path = self.path.with_extension("tmp");
        {
            let mut tmp = File::create(&tmp_path)?;
            for entry in &uncommitted {
                let line = serde_json::to_string(entry).map_err(|e| {
                    std::io::Error::new(std::io::ErrorKind::Other, e)
                })?;
                writeln!(tmp, "{}", line)?;
            }
            tmp.sync_all()?;
        }

        // Swap files
        fs::rename(&tmp_path, &self.path)?;

        // Reopen the file handle for future appends
        let new_file = OpenOptions::new()
            .create(true)
            .append(true)
            .open(&self.path)?;

        let mut file = self.file.lock().map_err(|_| {
            std::io::Error::new(std::io::ErrorKind::Other, "WAL lock poisoned")
        })?;
        *file = new_file;

        // Clear committed/failed entries from index
        let to_remove: Vec<String> = self.entries_index
            .iter()
            .filter(|e| e.status == WalStatus::Committed || e.status == WalStatus::Failed)
            .map(|e| e.key().clone())
            .collect();
        for key in to_remove {
            self.entries_index.remove(&key);
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;

    fn tmp_wal(name: &str) -> std::path::PathBuf {
        std::path::PathBuf::from(format!("/tmp/zc_wal_test_{}.jsonl", name))
    }

    fn cleanup(p: &std::path::Path) {
        let _ = fs::remove_file(p);
        let _ = fs::remove_file(p.with_extension("tmp"));
    }

    fn entry(request_id: &str, status: WalStatus) -> WalEntry {
        WalEntry {
            request_id: request_id.to_string(),
            user_id: "alice".to_string(),
            reservation_id: format!("res-{}", request_id),
            estimated_cost: 0.05,
            actual_cost: None,
            worker_id: "worker-1".to_string(),
            duration_ms: None,
            timestamp: Utc::now(),
            status,
        }
    }

    #[test]
    fn test_open_creates_file() {
        let path = tmp_wal("open");
        cleanup(&path);
        let _wal = Wal::open(&path).unwrap();
        assert!(path.exists());
        cleanup(&path);
    }

    #[test]
    fn test_append_then_flush_persists_to_disk() {
        let path = tmp_wal("append");
        cleanup(&path);

        let wal = Wal::open(&path).unwrap();
        wal.append(&entry("req-1", WalStatus::Reserved)).unwrap();
        wal.flush_buffer().unwrap();

        let disk = read_all_from_disk(&path).unwrap();
        assert_eq!(disk.len(), 1);
        assert_eq!(disk[0].request_id, "req-1");
        assert_eq!(disk[0].status, WalStatus::Reserved);

        cleanup(&path);
    }

    #[test]
    fn test_update_status_appends_new_line() {
        let path = tmp_wal("update");
        cleanup(&path);

        let wal = Wal::open(&path).unwrap();
        wal.append(&entry("req-2", WalStatus::Reserved)).unwrap();
        wal.update_status("req-2", WalStatus::Committed, Some(0.03), Some(100.0)).unwrap();
        wal.flush_buffer().unwrap();

        // Two lines on disk: reserved + committed
        let disk = read_all_from_disk(&path).unwrap();
        assert_eq!(disk.len(), 2);
        assert_eq!(disk[1].status, WalStatus::Committed);
        assert_eq!(disk[1].actual_cost, Some(0.03));

        cleanup(&path);
    }

    #[test]
    fn test_update_status_unknown_request_is_noop() {
        let path = tmp_wal("update_noop");
        cleanup(&path);
        let wal = Wal::open(&path).unwrap();
        // Should not panic or error
        wal.update_status("nonexistent", WalStatus::Committed, None, None).unwrap();
        wal.flush_buffer().unwrap();
        let disk = read_all_from_disk(&path).unwrap();
        assert!(disk.is_empty());
        cleanup(&path);
    }

    #[test]
    fn test_read_uncommitted_excludes_committed_and_failed() {
        let path = tmp_wal("uncommitted");
        cleanup(&path);

        let wal = Wal::open(&path).unwrap();
        wal.append(&entry("req-a", WalStatus::Reserved)).unwrap();
        wal.append(&entry("req-b", WalStatus::Reserved)).unwrap();
        wal.append(&entry("req-c", WalStatus::Reserved)).unwrap();
        wal.update_status("req-b", WalStatus::Committed, Some(0.01), Some(10.0)).unwrap();
        wal.update_status("req-c", WalStatus::Failed, None, None).unwrap();
        wal.flush_buffer().unwrap();

        let uncommitted = wal.read_uncommitted().unwrap();
        let ids: Vec<&str> = uncommitted.iter().map(|e| e.request_id.as_str()).collect();
        assert!(ids.contains(&"req-a"), "req-a should be uncommitted");
        assert!(!ids.contains(&"req-b"), "req-b (committed) should be excluded");
        assert!(!ids.contains(&"req-c"), "req-c (failed) should be excluded");

        cleanup(&path);
    }

    #[test]
    fn test_compact_removes_committed_entries_from_disk() {
        let path = tmp_wal("compact");
        cleanup(&path);

        let wal = Wal::open(&path).unwrap();
        wal.append(&entry("req-keep", WalStatus::Reserved)).unwrap();
        wal.append(&entry("req-done", WalStatus::Reserved)).unwrap();
        wal.update_status("req-done", WalStatus::Committed, Some(0.01), Some(5.0)).unwrap();
        wal.flush_buffer().unwrap();

        wal.compact().unwrap();

        let disk = read_all_from_disk(&path).unwrap();
        let ids: Vec<&str> = disk.iter().map(|e| e.request_id.as_str()).collect();
        assert!(ids.contains(&"req-keep"), "uncommitted entry should survive compaction");
        assert!(!ids.contains(&"req-done"), "committed entry should be removed");

        cleanup(&path);
    }

    #[test]
    fn test_compact_also_clears_failed_entries() {
        let path = tmp_wal("compact_failed");
        cleanup(&path);

        let wal = Wal::open(&path).unwrap();
        wal.append(&entry("req-fail", WalStatus::Reserved)).unwrap();
        wal.update_status("req-fail", WalStatus::Failed, None, None).unwrap();
        wal.flush_buffer().unwrap();
        wal.compact().unwrap();

        let disk = read_all_from_disk(&path).unwrap();
        assert!(disk.is_empty());

        cleanup(&path);
    }

    #[test]
    fn test_reload_populates_index_from_disk() {
        let path = tmp_wal("reload");
        cleanup(&path);

        // Write and close
        {
            let wal = Wal::open(&path).unwrap();
            wal.append(&entry("req-persist", WalStatus::Reserved)).unwrap();
            wal.flush_buffer().unwrap();
        }

        // Reopen: index must be repopulated from disk
        {
            let wal = Wal::open(&path).unwrap();
            let uncommitted = wal.read_uncommitted().unwrap();
            assert!(
                uncommitted.iter().any(|e| e.request_id == "req-persist"),
                "Reloaded WAL must surface persisted reserved entry"
            );
        }

        cleanup(&path);
    }

    #[test]
    fn test_full_lifecycle_reserved_executed_committed() {
        let path = tmp_wal("lifecycle");
        cleanup(&path);

        let wal = Wal::open(&path).unwrap();
        let e = entry("req-life", WalStatus::Reserved);
        wal.append(&e).unwrap();

        // Transition to Executed
        wal.update_status("req-life", WalStatus::Executed, None, Some(250.0)).unwrap();

        // Transition to Committed
        wal.update_status("req-life", WalStatus::Committed, Some(0.04), Some(250.0)).unwrap();
        wal.flush_buffer().unwrap();

        // No uncommitted entries should remain
        let uncommitted = wal.read_uncommitted().unwrap();
        assert!(!uncommitted.iter().any(|e| e.request_id == "req-life"));

        cleanup(&path);
    }
}