//! Log compaction and snapshots.

use serde::{Deserialize, Serialize};

/// Metadata for a snapshot.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct SnapshotMetadata {
    /// Index of last entry included in snapshot.
    pub last_included_index: u64,
    /// Term of last entry included in snapshot.
    pub last_included_term: u64,
    /// Cluster configuration at snapshot time.
    pub config: SnapshotConfig,
    /// Size of snapshot data in bytes.
    pub size: u64,
    /// Timestamp when snapshot was created.
    pub created_at: u64,
}

/// Cluster configuration stored in snapshot.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct SnapshotConfig {
    /// Node IDs in the cluster.
    pub nodes: Vec<u64>,
}

impl Default for SnapshotConfig {
    fn default() -> Self {
        Self { nodes: Vec::new() }
    }
}

/// A complete snapshot of the state machine.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Snapshot {
    /// Snapshot metadata.
    pub metadata: SnapshotMetadata,
    /// Snapshot data (serialized state machine).
    pub data: Vec<u8>,
}

impl Snapshot {
    /// Creates a new snapshot.
    pub fn new(
        last_included_index: u64,
        last_included_term: u64,
        config: SnapshotConfig,
        data: Vec<u8>,
    ) -> Self {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_millis() as u64;

        Self {
            metadata: SnapshotMetadata {
                last_included_index,
                last_included_term,
                config,
                size: data.len() as u64,
                created_at: now,
            },
            data,
        }
    }

    /// Serializes the snapshot to bytes.
    pub fn to_bytes(&self) -> Vec<u8> {
        bincode::serialize(self).unwrap_or_default()
    }

    /// Deserializes from bytes.
    pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
        bincode::deserialize(bytes).ok()
    }
}

/// Manages snapshot creation and storage.
pub struct SnapshotManager {
    /// Threshold for log entries before snapshotting.
    snapshot_threshold: u64,
    /// Current snapshot (if any).
    current_snapshot: Option<Snapshot>,
    /// Pending snapshot being received from leader.
    pending_snapshot: Option<PendingSnapshot>,
}

/// Snapshot being received in chunks.
struct PendingSnapshot {
    metadata: SnapshotMetadata,
    chunks: Vec<Vec<u8>>,
    expected_offset: u64,
}

impl SnapshotManager {
    /// Creates a new snapshot manager.
    pub fn new(snapshot_threshold: u64) -> Self {
        Self {
            snapshot_threshold,
            current_snapshot: None,
            pending_snapshot: None,
        }
    }

    /// Returns the threshold for triggering snapshots.
    pub fn threshold(&self) -> u64 {
        self.snapshot_threshold
    }

    /// Checks if a snapshot should be taken.
    pub fn should_snapshot(&self, log_size: u64, last_snapshot_index: u64) -> bool {
        log_size - last_snapshot_index >= self.snapshot_threshold
    }

    /// Gets the current snapshot.
    pub fn get_snapshot(&self) -> Option<&Snapshot> {
        self.current_snapshot.as_ref()
    }

    /// Gets the last included index of the current snapshot.
    pub fn last_included_index(&self) -> u64 {
        self.current_snapshot
            .as_ref()
            .map(|s| s.metadata.last_included_index)
            .unwrap_or(0)
    }

    /// Gets the last included term of the current snapshot.
    pub fn last_included_term(&self) -> u64 {
        self.current_snapshot
            .as_ref()
            .map(|s| s.metadata.last_included_term)
            .unwrap_or(0)
    }

    /// Creates a new snapshot.
    pub fn create_snapshot(
        &mut self,
        last_included_index: u64,
        last_included_term: u64,
        config: SnapshotConfig,
        data: Vec<u8>,
    ) {
        let snapshot = Snapshot::new(last_included_index, last_included_term, config, data);
        self.current_snapshot = Some(snapshot);
    }

    /// Starts receiving a snapshot from leader.
    pub fn start_receiving(
        &mut self,
        last_included_index: u64,
        last_included_term: u64,
        offset: u64,
        data: Vec<u8>,
    ) -> bool {
        if offset != 0 {
            // First chunk should be at offset 0
            return false;
        }

        self.pending_snapshot = Some(PendingSnapshot {
            metadata: SnapshotMetadata {
                last_included_index,
                last_included_term,
                config: SnapshotConfig::default(),
                size: 0,
                created_at: 0,
            },
            chunks: vec![data],
            expected_offset: 0,
        });

        true
    }

    /// Adds a chunk to the pending snapshot.
    pub fn add_chunk(&mut self, offset: u64, data: Vec<u8>) -> bool {
        if let Some(ref mut pending) = self.pending_snapshot {
            if offset == pending.expected_offset + pending.chunks.iter().map(|c| c.len() as u64).sum::<u64>() {
                pending.chunks.push(data);
                return true;
            }
        }
        false
    }

    /// Finalizes the pending snapshot.
    pub fn finalize_snapshot(&mut self) -> Option<Snapshot> {
        if let Some(pending) = self.pending_snapshot.take() {
            let data: Vec<u8> = pending.chunks.into_iter().flatten().collect();
            let snapshot = Snapshot::new(
                pending.metadata.last_included_index,
                pending.metadata.last_included_term,
                pending.metadata.config,
                data,
            );
            self.current_snapshot = Some(snapshot.clone());
            return Some(snapshot);
        }
        None
    }

    /// Cancels receiving a snapshot.
    pub fn cancel_receiving(&mut self) {
        self.pending_snapshot = None;
    }

    /// Splits snapshot into chunks for transmission.
    pub fn chunk_snapshot(&self, chunk_size: usize) -> Vec<(u64, Vec<u8>, bool)> {
        let mut chunks = Vec::new();

        if let Some(ref snapshot) = self.current_snapshot {
            let data = &snapshot.data;
            let mut offset = 0;

            while offset < data.len() {
                let end = (offset + chunk_size).min(data.len());
                let chunk = data[offset..end].to_vec();
                let done = end == data.len();
                chunks.push((offset as u64, chunk, done));
                offset = end;
            }
        }

        chunks
    }
}

impl Default for SnapshotManager {
    fn default() -> Self {
        Self::new(10000) // Default: snapshot every 10k entries
    }
}

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

    #[test]
    fn test_snapshot_creation() {
        let snapshot = Snapshot::new(100, 5, SnapshotConfig::default(), vec![1, 2, 3, 4, 5]);

        assert_eq!(snapshot.metadata.last_included_index, 100);
        assert_eq!(snapshot.metadata.last_included_term, 5);
        assert_eq!(snapshot.metadata.size, 5);
    }

    #[test]
    fn test_snapshot_serialization() {
        let snapshot = Snapshot::new(100, 5, SnapshotConfig::default(), vec![1, 2, 3, 4, 5]);
        let bytes = snapshot.to_bytes();
        let decoded = Snapshot::from_bytes(&bytes).unwrap();

        assert_eq!(decoded.metadata.last_included_index, 100);
        assert_eq!(decoded.data, vec![1, 2, 3, 4, 5]);
    }

    #[test]
    fn test_snapshot_manager() {
        let mut manager = SnapshotManager::new(100);

        assert!(manager.should_snapshot(150, 0));
        assert!(!manager.should_snapshot(50, 0));

        manager.create_snapshot(100, 5, SnapshotConfig::default(), vec![1, 2, 3]);

        assert_eq!(manager.last_included_index(), 100);
        assert_eq!(manager.last_included_term(), 5);
    }

    #[test]
    fn test_chunking() {
        let mut manager = SnapshotManager::new(100);
        manager.create_snapshot(100, 5, SnapshotConfig::default(), vec![0; 250]);

        let chunks = manager.chunk_snapshot(100);

        assert_eq!(chunks.len(), 3);
        assert_eq!(chunks[0].0, 0);
        assert_eq!(chunks[0].1.len(), 100);
        assert!(!chunks[0].2); // not done
        assert_eq!(chunks[2].0, 200);
        assert_eq!(chunks[2].1.len(), 50);
        assert!(chunks[2].2); // done
    }

    #[test]
    fn test_receiving_snapshot() {
        let mut manager = SnapshotManager::new(100);

        // Start receiving
        assert!(manager.start_receiving(100, 5, 0, vec![1, 2, 3]));

        // Add more chunks
        assert!(manager.add_chunk(3, vec![4, 5, 6]));

        // Finalize
        let snapshot = manager.finalize_snapshot().unwrap();
        assert_eq!(snapshot.data, vec![1, 2, 3, 4, 5, 6]);
    }
}
