use std::fs::{self, File};
use std::io::Read;
use std::path::Path;

use super::checkpoint::CheckpointManager;
use super::error::WalError;
use super::format::{self, BatchHeader, EventRecord, HEADER_SIZE, MAGIC};
use super::segment::list_segments;

/// Result of WAL recovery: replayed events and the next usable sequence number.
pub struct RecoveryResult {
    /// Events replayed from the WAL since the last checkpoint.
    pub events: Vec<EventRecord>,
    /// The next sequence number to assign (one past the last valid event).
    pub next_seq: u64,
}

/// Recover the WAL from disk.
///
/// Reads the checkpoint (if any), then scans all segment files from the
/// checkpoint position forward. Validates each batch with two-phase
/// checking (magic + bounds, then BLAKE3). Truncates any corrupted tail.
///
/// Returns the replayed events and the next sequence number to assign.
///
/// # Errors
///
/// Returns `WalError::Io` on filesystem failure, or `WalError::Corruption`
/// if a segment is corrupted in a way that cannot be recovered by truncation.
pub fn recover(dir: &Path) -> Result<RecoveryResult, WalError> {
    let checkpoint = CheckpointManager::read(dir)?;
    let checkpoint_seq = checkpoint.map_or(0, |(seq, _)| seq);

    let segments = list_segments(dir)?;
    if segments.is_empty() {
        return Ok(RecoveryResult {
            events: Vec::new(),
            // New events must have seq > checkpoint_seq so WAL replay (which uses
            // strict greater-than) never skips them. Start one past the checkpoint.
            // max(1) ensures seq=0 (the dedup sentinel) is never assigned.
            next_seq: checkpoint_seq.saturating_add(1).max(1),
        });
    }

    let mut all_events = Vec::new();
    // Same invariant for the loop: new events must have seq > checkpoint_seq.
    // If all segment events are at or below checkpoint_seq, this floor ensures
    // the writer still starts past the checkpoint boundary.
    let mut next_seq = checkpoint_seq.saturating_add(1).max(1);

    for (seg_first_seq, seg_path) in &segments {
        // Skip segments that are entirely before the checkpoint.
        // We need to scan segments that *might* contain post-checkpoint events.
        // A segment starting at seg_first_seq could contain events up to some
        // higher sequence number, so we only skip if we can definitively
        // determine all events are before the checkpoint.
        // Since we scan forward and track next_seq, we handle this during event iteration.

        let scan_result = scan_segment(seg_path)?;
        for (header, events) in scan_result {
            for (i, event) in events.into_iter().enumerate() {
                let event_seq = header.first_seq + i as u64;
                // Only replay events strictly after the checkpoint boundary.
                // checkpoint_seq is the last sequence that was materialized
                // to durable storage — it must not be replayed again.
                if event_seq > checkpoint_seq {
                    all_events.push(event);
                }
                // Track the highest sequence number seen
                let candidate = event_seq + 1;
                if candidate > next_seq {
                    next_seq = candidate;
                }
            }
        }

        // If we're past the checkpoint and the segment might have been partially
        // written, we already handled truncation in scan_segment.
        let _ = seg_first_seq; // suppress unused warning in the skip logic comment
    }

    Ok(RecoveryResult {
        events: all_events,
        next_seq,
    })
}

/// Scan all WAL segment files and return every event record.
///
/// Unlike [`recover`], this reads ALL events regardless of the checkpoint
/// boundary. Used for RLHF export and offline analysis.
///
/// # Errors
///
/// Returns `WalError::Io` if a segment file cannot be opened.
pub fn read_all_events(dir: &Path) -> Result<Vec<EventRecord>, WalError> {
    let segments = list_segments(dir)?;
    let mut all_events = Vec::new();
    for (_seq, path) in &segments {
        for (_header, events) in scan_segment(path)? {
            all_events.extend(events);
        }
    }
    Ok(all_events)
}

/// Scan a single segment file, returning all valid batches.
///
/// On encountering a corrupted or truncated batch, truncates the file
/// to the last valid offset and stops. This handles torn writes from crashes.
pub(crate) fn scan_segment(path: &Path) -> Result<Vec<(BatchHeader, Vec<EventRecord>)>, WalError> {
    let mut file = File::open(path)?;
    let file_len = file.metadata()?.len();

    let mut data = Vec::new();
    file.read_to_end(&mut data)?;
    drop(file);

    let mut batches = Vec::new();
    let mut offset: usize = 0;
    let mut last_valid_offset: usize = 0;

    while offset < data.len() {
        // Phase 1: Can we read a header?
        if data.len() - offset < HEADER_SIZE {
            tracing::warn!(
                path = %path.display(),
                offset,
                remaining = data.len() - offset,
                "truncated header at end of segment"
            );
            break;
        }

        // Phase 1: Check magic bytes
        if data[offset..offset + 4] != MAGIC {
            tracing::warn!(
                path = %path.display(),
                offset,
                "invalid magic bytes, stopping scan"
            );
            break;
        }

        // Phase 1: Read payload_len to check bounds
        let payload_len =
            u32::from_le_bytes(data[offset + 24..offset + 28].try_into().map_err(|_| {
                WalError::Corruption {
                    message: "failed to read payload_len".into(),
                }
            })?) as usize;

        let batch_end = offset + HEADER_SIZE + payload_len;
        if batch_end > data.len() {
            tracing::warn!(
                path = %path.display(),
                offset,
                payload_len,
                file_len,
                "truncated payload, stopping scan"
            );
            break;
        }

        // Phase 2: Full decode with BLAKE3 verification
        match format::decode_batch(&data[offset..batch_end]) {
            Ok((header, events)) => {
                batches.push((header, events));
                last_valid_offset = batch_end;
                offset = batch_end;
            }
            Err(e) => {
                tracing::warn!(
                    path = %path.display(),
                    offset,
                    error = %e,
                    "batch decode failed, stopping scan"
                );
                break;
            }
        }
    }

    // Truncate the file to remove any corrupted tail
    if last_valid_offset < data.len() {
        tracing::warn!(
            path = %path.display(),
            valid_bytes = last_valid_offset,
            total_bytes = data.len(),
            "truncating corrupted tail"
        );
        let file = fs::OpenOptions::new().write(true).open(path)?;
        file.set_len(last_valid_offset as u64)?;
        file.sync_all()?;
    }

    Ok(batches)
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::super::format::{EventRecord, encode_batch};
    use super::*;
    use crate::replication::ShardId;

    fn sample_event(id: u64, ts: u64) -> EventRecord {
        EventRecord {
            entity_id: id,
            signal_type: 1,
            weight: 1.0,
            timestamp_nanos: ts,
        }
    }

    #[test]
    fn recover_empty_dir() {
        let dir = tempfile::tempdir().expect("tempdir creation should succeed");
        let result = recover(dir.path()).expect("recover should succeed");
        assert!(result.events.is_empty());
        // No checkpoint, no segments: next_seq starts at 1 (seq=0 is the dedup sentinel).
        assert_eq!(result.next_seq, 1);
    }

    #[test]
    fn recover_single_batch() {
        let dir = tempfile::tempdir().expect("tempdir creation should succeed");
        let events = vec![sample_event(1, 1000), sample_event(2, 2000)];
        let batch_bytes = encode_batch(&events, 1, 1000).expect("encode should succeed");

        let seg_name = super::super::segment::segment_filename(ShardId::SINGLE, 1);
        fs::write(dir.path().join(seg_name), &batch_bytes).expect("write should succeed");

        let result = recover(dir.path()).expect("recover should succeed");
        assert_eq!(result.events.len(), 2);
        assert_eq!(result.next_seq, 3); // first_seq=1, 2 events -> next=3
    }

    #[test]
    fn recover_multiple_batches() {
        let dir = tempfile::tempdir().expect("tempdir creation should succeed");

        let events1 = vec![sample_event(1, 1000)];
        let events2 = vec![sample_event(2, 2000)];

        let batch1 = encode_batch(&events1, 1, 1000).expect("encode should succeed");
        let batch2 = encode_batch(&events2, 2, 2000).expect("encode should succeed");

        let mut segment_data = batch1;
        segment_data.extend_from_slice(&batch2);

        let seg_name = super::super::segment::segment_filename(ShardId::SINGLE, 1);
        fs::write(dir.path().join(seg_name), &segment_data).expect("write should succeed");

        let result = recover(dir.path()).expect("recover should succeed");
        assert_eq!(result.events.len(), 2);
        assert_eq!(result.next_seq, 3);
    }

    #[test]
    fn recover_truncated_tail() {
        let dir = tempfile::tempdir().expect("tempdir creation should succeed");

        let events = vec![sample_event(1, 1000)];
        let batch_bytes = encode_batch(&events, 1, 1000).expect("encode should succeed");

        // Write a valid batch followed by garbage (simulating a torn write)
        let mut data = batch_bytes.clone();
        data.extend_from_slice(&[0xDE, 0xAD, 0xBE, 0xEF, 0x00, 0x11]);

        let seg_name = super::super::segment::segment_filename(ShardId::SINGLE, 1);
        let seg_path = dir.path().join(&seg_name);
        fs::write(&seg_path, &data).expect("write should succeed");

        let result = recover(dir.path()).expect("recover should succeed");
        assert_eq!(result.events.len(), 1);
        assert_eq!(result.next_seq, 2);

        // Verify file was truncated to valid length
        let file_len = fs::metadata(&seg_path)
            .expect("metadata should succeed")
            .len();
        assert_eq!(file_len, batch_bytes.len() as u64);
    }

    #[test]
    fn recover_with_checkpoint_skips_old_events() {
        let dir = tempfile::tempdir().expect("tempdir creation should succeed");

        // Write a checkpoint at seq=4 (events through seq=4 are materialized; 5 and 6 need replay)
        CheckpointManager::write(dir.path(), 4, 4000).expect("write should succeed");

        // Batch 1: events at seq 1..=3
        let events1: Vec<EventRecord> = (1..=3).map(|i| sample_event(i, i * 1000)).collect();
        let batch1 = encode_batch(&events1, 1, 1000).expect("encode should succeed");

        // Batch 2: events at seq 4..=6
        let events2: Vec<EventRecord> = (4..=6).map(|i| sample_event(i, i * 1000)).collect();
        let batch2 = encode_batch(&events2, 4, 4000).expect("encode should succeed");

        let mut segment_data = batch1;
        segment_data.extend_from_slice(&batch2);

        let seg_name = super::super::segment::segment_filename(ShardId::SINGLE, 1);
        fs::write(dir.path().join(seg_name), &segment_data).expect("write should succeed");

        let result = recover(dir.path()).expect("recover should succeed");
        // Events at seq 5 and 6 should be replayed (> checkpoint_seq=4)
        assert_eq!(result.events.len(), 2);
        assert_eq!(result.events[0].entity_id, 5);
        assert_eq!(result.events[1].entity_id, 6);
        assert_eq!(result.next_seq, 7);
    }

    #[test]
    fn recover_corrupted_batch_in_middle() {
        let dir = tempfile::tempdir().expect("tempdir creation should succeed");

        let events1 = vec![sample_event(1, 1000)];
        let events2 = vec![sample_event(2, 2000)];

        let batch1 = encode_batch(&events1, 1, 1000).expect("encode should succeed");
        let mut batch2 = encode_batch(&events2, 2, 2000).expect("encode should succeed");
        // Corrupt batch2's payload
        batch2[HEADER_SIZE + 5] ^= 0xFF;

        let mut segment_data = batch1;
        segment_data.extend_from_slice(&batch2);

        let seg_name = super::super::segment::segment_filename(ShardId::SINGLE, 1);
        fs::write(dir.path().join(seg_name), &segment_data).expect("write should succeed");

        let result = recover(dir.path()).expect("recover should succeed");
        // Only batch1 is valid
        assert_eq!(result.events.len(), 1);
        assert_eq!(result.next_seq, 2);
    }

    #[test]
    fn recover_empty_segment_file() {
        let dir = tempfile::tempdir().expect("tempdir creation should succeed");
        let seg_name = super::super::segment::segment_filename(ShardId::SINGLE, 1);
        fs::write(dir.path().join(seg_name), []).expect("write should succeed");

        let result = recover(dir.path()).expect("recover should succeed");
        assert!(result.events.is_empty());
        // No checkpoint, empty segment: next_seq starts at 1 (seq=0 is the dedup sentinel).
        assert_eq!(result.next_seq, 1);
    }

    #[test]
    fn recover_multiple_segments() {
        let dir = tempfile::tempdir().expect("tempdir creation should succeed");

        let events1 = vec![sample_event(1, 1000)];
        let events2 = vec![sample_event(2, 2000)];

        let batch1 = encode_batch(&events1, 1, 1000).expect("encode should succeed");
        let batch2 = encode_batch(&events2, 2, 2000).expect("encode should succeed");

        let seg1 = super::super::segment::segment_filename(ShardId::SINGLE, 1);
        let seg2 = super::super::segment::segment_filename(ShardId::SINGLE, 2);
        fs::write(dir.path().join(seg1), &batch1).expect("write should succeed");
        fs::write(dir.path().join(seg2), &batch2).expect("write should succeed");

        let result = recover(dir.path()).expect("recover should succeed");
        assert_eq!(result.events.len(), 2);
        assert_eq!(result.next_seq, 3);
    }
}
