lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0
//
// Replication Stream Protocol (Send/Receive)
// ZFS-style incremental replication for backup and migration.

use alloc::format;
use alloc::string::String;
use alloc::vec::Vec;

use crate::fscore::structs::{Blkptr, DnodePhys};
use crate::integrity::checksum::Checksum;
use crate::io::pipeline::Pipeline;

/// Send stream record types for replication protocol
#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum SendRecordType {
    /// Stream begin marker
    Begin = 1, // Stream begin marker
    /// Stream end marker
    End = 2, // Stream end marker
    /// New object/file creation
    Object = 3, // New object/file
    /// Free object range operation
    Freeobjects = 4, // Free object range
    /// Data block write
    Write = 5, // Data block
    /// Free block range operation
    Free = 6, // Free block range
    /// Free specific block pointer
    FreeBlockRef = 7, // Free specific block pointer
    /// Dataset properties update
    Properties = 8, // Dataset properties
    /// Snapshot creation marker
    Snapshot = 9, // Snapshot marker
    /// Destroy object operation
    Destroy = 10, // Destroy object
}

/// Send stream header (64 bytes)
#[repr(C, packed)]
#[derive(Debug, Clone, Copy)]
pub struct SendStreamHeader {
    /// Magic number (0x53454E44 = "SEND")
    pub magic: u64, // 0x53454E44 ("SEND")
    /// Protocol version number
    pub version: u32, // Protocol version
    /// Stream flags
    pub flags: u32, // Stream flags
    /// Source dataset GUID
    pub from_guid: u64, // Source dataset GUID
    /// Target dataset GUID (0 for full send)
    pub to_guid: u64, // Target dataset GUID (0 for full send)
    /// Starting transaction group number
    pub from_txg: u64, // Starting TXG
    /// Ending transaction group number
    pub to_txg: u64, // Ending TXG
    /// Stream creation timestamp
    pub creation_time: u64, // Stream creation timestamp
    /// Fletcher4 checksum of entire stream
    pub checksum: u64, // Fletcher4 of entire stream
    /// Reserved for future use
    pub reserved: [u64; 3],
}

/// Send stream magic number: "SEND"
const SEND_MAGIC: u64 = 0x53454E44; // "SEND"

impl SendStreamHeader {
    /// Create a new send stream header
    pub fn new(from_guid: u64, to_guid: u64, from_txg: u64, to_txg: u64) -> Self {
        Self {
            magic: SEND_MAGIC,
            version: 1,
            flags: 0,
            from_guid,
            to_guid,
            from_txg,
            to_txg,
            creation_time: 0, // Would use RDTSC
            checksum: 0,      // Calculated at stream end
            reserved: [0; 3],
        }
    }
}

/// Send stream record (variable length)
#[repr(C, packed)]
#[derive(Debug, Clone, Copy)]
pub struct SendRecord {
    /// Record type identifier
    pub record_type: u8,
    /// Record flags
    pub flags: u8,
    /// Payload data length
    pub length: u32, // Payload length
    /// Target object ID
    pub object_id: u64, // Target object
    /// Offset for write operations
    pub offset: u64, // Offset (for writes)
    /// Fletcher4 checksum of payload
    pub checksum: u64, // Fletcher4 of payload
}

impl SendRecord {
    /// Create a new send record
    pub fn new(record_type: SendRecordType, object_id: u64, offset: u64, length: u32) -> Self {
        Self {
            record_type: record_type as u8,
            flags: 0,
            length,
            object_id,
            offset,
            checksum: 0,
        }
    }
}

/// Send stream builder for creating replication streams
pub struct SendStream {
    /// Stream header with metadata
    pub header: SendStreamHeader,
    /// Serialized stream buffer
    pub buffer: Vec<u8>,
    /// Number of records in stream
    pub record_count: usize,
}

impl SendStream {
    /// Begin a new send stream
    pub fn begin_full_send(dataset_guid: u64, to_txg: u64) -> Self {
        crate::lcpfs_println!(
            "[ SEND ] Beginning full send for GUID {:x} @ TXG {}",
            dataset_guid,
            to_txg
        );

        let header = SendStreamHeader::new(dataset_guid, 0, 0, to_txg);

        let mut stream = Self {
            header,
            buffer: Vec::new(),
            record_count: 0,
        };

        // Write header
        stream.write_header();

        stream
    }

    /// Begin incremental send (from snapshot to snapshot)
    pub fn begin_incremental_send(
        from_guid: u64,
        to_guid: u64,
        from_txg: u64,
        to_txg: u64,
    ) -> Self {
        crate::lcpfs_println!(
            "[ SEND ] Beginning incremental send {:x} -> {:x} (TXG {}-{})",
            from_guid,
            to_guid,
            from_txg,
            to_txg
        );

        let header = SendStreamHeader::new(from_guid, to_guid, from_txg, to_txg);

        let mut stream = Self {
            header,
            buffer: Vec::new(),
            record_count: 0,
        };

        stream.write_header();
        stream
    }

    /// Write stream header to buffer
    fn write_header(&mut self) {
        // SAFETY INVARIANTS:
        // 1. SendStreamHeader is #[repr(C, packed)] with stable layout
        // 2. self.header is valid, initialized SendStreamHeader
        // 3. All fields are primitive types (u64, u32, arrays - no Drop)
        // 4. Slice lifetime ≤ function scope (does not escape)
        // 5. size_of::<SendStreamHeader>() matches packed struct size
        //
        // VERIFICATION: TODO - Prove packed struct has no padding
        //
        // JUSTIFICATION:
        // ZFS send/receive protocol header (magic, version, flags).
        // Binary format for cross-system snapshot transfer.
        unsafe {
            let header_bytes = core::slice::from_raw_parts(
                &self.header as *const SendStreamHeader as *const u8,
                core::mem::size_of::<SendStreamHeader>(),
            );
            self.buffer.extend_from_slice(header_bytes);
        }
    }

    /// Add a data block write record
    pub fn write_block(&mut self, object_id: u64, offset: u64, data: &[u8]) {
        let record = SendRecord::new(SendRecordType::Write, object_id, offset, data.len() as u32);

        self.write_record(&record, data);
    }

    /// Add object creation record
    pub fn write_object(&mut self, object_id: u64, object_type: u8, dnode: &DnodePhys) {
        let record = SendRecord::new(
            SendRecordType::Object,
            object_id,
            0,
            core::mem::size_of::<DnodePhys>() as u32,
        );

        // Serialize dnode as payload
        // SAFETY INVARIANTS:
        // 1. DnodePhys is #[repr(C)] with stable layout
        // 2. dnode reference is valid and properly initialized
        // 3. All fields are primitive types or arrays (no Drop)
        // 4. Slice lifetime ≤ function scope (does not escape)
        // 5. size_of::<DnodePhys>() matches struct size
        //
        // VERIFICATION: TODO - Same as lcpfs_dmu.rs (DnodePhys stability)
        //
        // JUSTIFICATION:
        // Send stream includes object metadata as serialized DnodePhys.
        // Required for recreating objects during receive operation.
        let payload = unsafe {
            core::slice::from_raw_parts(
                dnode as *const DnodePhys as *const u8,
                core::mem::size_of::<DnodePhys>(),
            )
        };

        self.write_record(&record, payload);
    }

    /// Add free block record
    pub fn write_free(&mut self, object_id: u64, offset: u64, length: u64) {
        let record = SendRecord::new(SendRecordType::Free, object_id, offset, length as u32);

        self.write_record(&record, &[]);
    }

    /// Add snapshot marker
    pub fn write_snapshot(&mut self, snap_name: &str, guid: u64) {
        let record = SendRecord::new(SendRecordType::Snapshot, guid, 0, snap_name.len() as u32);

        self.write_record(&record, snap_name.as_bytes());
    }

    /// Write a record to the stream
    fn write_record(&mut self, record: &SendRecord, payload: &[u8]) {
        // Calculate checksum
        let mut record_with_checksum = *record;
        record_with_checksum.checksum = Checksum::calculate(payload).first();

        // Serialize record header
        // SAFETY INVARIANTS:
        // 1. SendRecord is #[repr(C, packed)] with stable layout
        // 2. record_with_checksum is valid, initialized SendRecord on stack
        // 3. All fields are primitive types (u8, u32, u64 - no Drop)
        // 4. Slice lifetime ≤ function scope
        // 5. Packed struct has no padding bytes
        //
        // VERIFICATION: TODO - Prove SendRecord packed layout deterministic
        //
        // JUSTIFICATION:
        // Each send stream record (write/free/snapshot) has binary header.
        // Serialization required for ZFS send/receive wire protocol.
        unsafe {
            let record_bytes = core::slice::from_raw_parts(
                &record_with_checksum as *const SendRecord as *const u8,
                core::mem::size_of::<SendRecord>(),
            );
            self.buffer.extend_from_slice(record_bytes);
        }

        // Append payload
        self.buffer.extend_from_slice(payload);

        self.record_count += 1;
    }

    /// Finalize stream and return serialized data
    pub fn finalize(mut self) -> Vec<u8> {
        // Add stream end marker
        let end_record = SendRecord::new(SendRecordType::End, 0, 0, 0);
        self.write_record(&end_record, &[]);

        // Calculate final checksum of entire stream (except header checksum field)
        let stream_checksum = Checksum::calculate(&self.buffer).first();

        // Update header checksum
        let header_offset = offset_of_checksum_in_header();
        // SAFETY INVARIANTS:
        // 1. header_offset points to checksum field in SendStreamHeader
        // 2. buffer.len() ≥ header_offset + 8 (header written by write_header)
        // 3. Writing u64 at header_offset (field is properly aligned)
        // 4. Mutable ownership of buffer (no aliasing)
        // 5. write_unaligned handles potential misalignment in Vec
        //
        // VERIFICATION: TODO - Prove offset_of_checksum_in_header() correctness
        //
        // JUSTIFICATION:
        // Stream checksum stored in header for integrity verification.
        // In-place update of checksum field after stream construction.
        unsafe {
            let checksum_ptr = self.buffer.as_mut_ptr().add(header_offset) as *mut u64;
            *checksum_ptr = stream_checksum;
        }

        crate::lcpfs_println!(
            "[ SEND ] Finalized stream: {} records, {} bytes",
            self.record_count,
            self.buffer.len()
        );

        self.buffer
    }

    /// Get stream size
    pub fn size(&self) -> usize {
        self.buffer.len()
    }
}

/// Calculate offset of checksum field in header (safe method)
fn offset_of_checksum_in_header() -> usize {
    // Use size_of instead of manual calculation to be future-proof
    // Offset = size of all fields before checksum
    use core::mem::size_of;

    size_of::<u64>() +      // magic
    size_of::<u32>() +      // version
    size_of::<u32>() +      // flags
    size_of::<u64>() +      // from_guid
    size_of::<u64>() +      // to_guid
    size_of::<u64>() +      // from_txg
    size_of::<u64>() +      // to_txg
    size_of::<u64>() // creation_time
    // Next field is checksum
}

/// Receive stream parser for applying replication streams
pub struct ReceiveStream {
    /// Stream data buffer
    pub data: Vec<u8>,
    /// Current read offset in stream
    pub offset: usize,
    /// Parsed stream header
    pub header: Option<SendStreamHeader>,
    /// Number of records processed
    pub processed_records: usize,
    /// Whether stream completed with End marker
    stream_complete: bool,
}

impl ReceiveStream {
    /// Begin receiving a stream
    pub fn begin_receive(data: Vec<u8>) -> Result<Self, &'static str> {
        let mut stream = Self {
            data,
            offset: 0,
            header: None,
            processed_records: 0,
            stream_complete: false,
        };

        stream.read_header()?;
        Ok(stream)
    }

    /// Read and validate stream header
    fn read_header(&mut self) -> Result<(), &'static str> {
        if self.data.len() < core::mem::size_of::<SendStreamHeader>() {
            return Err("Stream too short");
        }

        // SAFETY INVARIANTS:
        // 1. data.len() ≥ size_of::<SendStreamHeader>() - checked above
        // 2. SendStreamHeader is #[repr(C, packed)] with stable layout
        // 3. read_unaligned handles packed struct (no alignment requirement)
        // 4. Data written by SendStream::write_header with same layout
        // 5. Owned SendStreamHeader returned (no lifetime issues)
        //
        // VERIFICATION: TODO - Prove send/receive format compatibility
        //
        // JUSTIFICATION:
        // Stream begins with SendStreamHeader for protocol negotiation.
        // Deserialization required to validate magic, version, GUIDs.
        let header: SendStreamHeader =
            unsafe { core::ptr::read_unaligned(self.data.as_ptr() as *const SendStreamHeader) };

        if header.magic != SEND_MAGIC {
            return Err("Invalid stream magic");
        }

        let from_guid = header.from_guid;
        let from_txg = header.from_txg;
        let to_txg = header.to_txg;
        crate::lcpfs_println!(
            "[ RECV ] Receiving stream: GUID {:x}, TXG {} -> {}",
            from_guid,
            from_txg,
            to_txg
        );

        self.offset = core::mem::size_of::<SendStreamHeader>();
        self.header = Some(header);

        Ok(())
    }

    /// Read next record
    pub fn read_next_record(&mut self) -> Result<Option<(SendRecord, Vec<u8>)>, &'static str> {
        if self.offset + core::mem::size_of::<SendRecord>() > self.data.len() {
            // Check if stream ended properly with an End marker
            if !self.stream_complete {
                return Err("Truncated stream: missing End marker");
            }
            return Ok(None);
        }

        // SAFETY INVARIANTS:
        // 1. offset + size_of::<SendRecord>() ≤ data.len() - checked above
        // 2. SendRecord is #[repr(C, packed)] with stable layout
        // 3. read_unaligned handles packed struct and potential misalignment
        // 4. Data written by SendStream::write_record with same layout
        // 5. Owned SendRecord returned (no lifetime issues)
        //
        // VERIFICATION: TODO - Prove record offset arithmetic correctness
        //
        // JUSTIFICATION:
        // Stream records (write/free/snapshot) have binary headers.
        // Sequential deserialization required for processing stream.
        let record: SendRecord = unsafe {
            core::ptr::read_unaligned(self.data.as_ptr().add(self.offset) as *const SendRecord)
        };

        self.offset += core::mem::size_of::<SendRecord>();

        // Check for end marker
        if record.record_type == SendRecordType::End as u8 {
            self.stream_complete = true;
            crate::lcpfs_println!(
                "[ RECV ] Stream complete: {} records processed",
                self.processed_records
            );
            return Ok(None);
        }

        // Read payload
        let payload_end = self.offset + record.length as usize;
        if payload_end > self.data.len() {
            return Err("Truncated payload");
        }

        let payload = self.data[self.offset..payload_end].to_vec();
        self.offset = payload_end;

        // Verify checksum
        let expected_checksum = Checksum::calculate(&payload).first();
        if record.checksum != expected_checksum {
            return Err("Payload checksum mismatch");
        }

        self.processed_records += 1;
        Ok(Some((record, payload)))
    }

    /// Process all records
    pub fn process_stream(&mut self) -> Result<(), &'static str> {
        while let Some((record, payload)) = self.read_next_record()? {
            self.apply_record(&record, &payload)?;
        }

        crate::lcpfs_println!(
            "[ RECV ] Successfully applied {} records",
            self.processed_records
        );
        Ok(())
    }

    /// Apply a single record
    fn apply_record(&self, record: &SendRecord, payload: &[u8]) -> Result<(), &'static str> {
        let record_type = match record.record_type {
            1 => SendRecordType::Begin,
            2 => SendRecordType::End,
            3 => SendRecordType::Object,
            4 => SendRecordType::Freeobjects,
            5 => SendRecordType::Write,
            6 => SendRecordType::Free,
            9 => SendRecordType::Snapshot,
            10 => SendRecordType::Destroy,
            _ => return Err("Unknown record type"),
        };

        match record_type {
            SendRecordType::Write => {
                let obj_id = record.object_id;
                let offset = record.offset;
                crate::lcpfs_println!(
                    "[ RECV ] Apply WRITE: obj={} offset={} len={}",
                    obj_id,
                    offset,
                    payload.len()
                );
                // In real implementation: write to DMU object at offset
            }
            SendRecordType::Object => {
                let obj_id = record.object_id;
                crate::lcpfs_println!("[ RECV ] Apply CREATE: obj={}", obj_id);
                // Create new DMU object
            }
            SendRecordType::Free => {
                let obj_id = record.object_id;
                let offset = record.offset;
                let length = record.length;
                crate::lcpfs_println!(
                    "[ RECV ] Apply FREE: obj={} offset={} len={}",
                    obj_id,
                    offset,
                    length
                );
                // Free block range
            }
            SendRecordType::Snapshot => {
                let snap_name = String::from_utf8_lossy(payload);
                crate::lcpfs_println!("[ RECV ] Apply SNAPSHOT: {}", snap_name);
            }
            _ => {}
        }

        Ok(())
    }
}

/// Example: Full dataset send
pub fn example_full_send(dataset_guid: u64, txg: u64) -> Vec<u8> {
    let mut stream = SendStream::begin_full_send(dataset_guid, txg);

    // Add some example data
    stream.write_object(1, 1, &DnodePhys::zero());
    stream.write_block(1, 0, b"Hello, LCPFS replication!");
    stream.write_block(1, 4096, b"More data here...");

    stream.finalize()
}

/// Example: Incremental send
pub fn example_incremental_send(
    from_guid: u64,
    to_guid: u64,
    from_txg: u64,
    to_txg: u64,
) -> Vec<u8> {
    let mut stream = SendStream::begin_incremental_send(from_guid, to_guid, from_txg, to_txg);

    // Only send changes between snapshots
    stream.write_block(1, 8192, b"Changed block");
    stream.write_free(2, 0, 4096);

    stream.finalize()
}