xet-client 1.6.0

use core::fmt;
use std::cmp::{Ordering, min};
use std::collections::{HashMap, HashSet};
use std::marker::PhantomData;
use std::str::FromStr;

use serde::{Deserialize, Serialize};
use serde_repr::{Deserialize_repr, Serialize_repr};
use thiserror::Error;
use xet_core_structures::merklehash::{MerkleHash, MerkleHashSubtree};

mod key;
pub use key::*;

/// Indicates a "session id" that clients can use to group together related requests
/// (e.g. all requests made to CAS to support a user-triggered upload (xorbs + shards)).
pub const SESSION_ID_HEADER: &str = "X-Xet-Session-Id";
/// Request id generated by CAS for a request.
pub const REQUEST_ID_HEADER: &str = "X-Request-Id";

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct UploadXorbResponse {
    pub was_inserted: bool,
}

/// These types are defined to help differentiate the Range<,> type aliases,
/// so that they don't silently cast to each other without range adjustments.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Default, Hash, Copy)]
pub struct _C;
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Default, Hash, Copy)]
pub struct _F;
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Default, Hash, Copy)]
pub struct _H;

/// Start and exclusive-end range for chunk content
pub type ChunkRange = Range<u32, _C>;
/// Start and exclusive-end range for file content
pub type FileRange = Range<u64, _F>;
/// Start and inclusive-end range for HTTP range content
pub type HttpRange = Range<u64, _H>;

impl FileRange {
    pub fn full() -> Self {
        Self::new(0, u64::MAX)
    }

    // consumes self and split the range into a segment of size `segment_size`
    // and a remainder.
    pub fn take_segment(self, segment_size: u64) -> (Self, Option<Self>) {
        let segment = FileRange {
            start: self.start,
            end: min(self.end, self.start + segment_size),
            _marker: PhantomData,
        };

        let remainder = if segment.end == self.end {
            None
        } else {
            Some(FileRange {
                start: segment.end,
                end: self.end,
                _marker: PhantomData,
            })
        };

        (segment, remainder)
    }

    pub fn length(&self) -> u64 {
        self.end - self.start
    }
}

impl From<HttpRange> for FileRange {
    fn from(value: HttpRange) -> Self {
        // right inclusive to right exclusive
        FileRange::new(value.start, value.end + 1)
    }
}

impl HttpRange {
    pub fn range_header(&self) -> String {
        format!("bytes={self}")
    }

    pub fn length(&self) -> u64 {
        self.end - self.start + 1
    }
}

impl From<FileRange> for HttpRange {
    fn from(value: FileRange) -> Self {
        // right exclusive to right inclusive
        HttpRange::new(value.start, value.end - 1)
    }
}

// note that the standard PartialOrd/Ord impls will first check `start` then `end`
#[derive(Serialize, Deserialize, Clone, Eq, PartialEq, PartialOrd, Ord, Default, Hash)]
pub struct Range<Idx, Kind> {
    pub start: Idx,
    pub end: Idx,
    #[serde(skip)]
    pub _marker: PhantomData<Kind>,
}

impl<Idx, _C> fmt::Debug for Range<Idx, _C>
where
    Idx: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Range")
            .field("start", &self.start)
            .field("end", &self.end)
            .finish()
    }
}

impl<Idx, Kind> Range<Idx, Kind> {
    pub fn new(start: Idx, end: Idx) -> Self {
        Self {
            start,
            end,
            _marker: PhantomData,
        }
    }
}

impl<T: Copy, Kind: Copy> Copy for Range<T, Kind> {}

impl<Idx: fmt::Display, Kind> fmt::Display for Range<Idx, Kind> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}-{}", self.start, self.end)
    }
}

#[derive(Error, Debug)]
pub enum RangeParseError<Idx: std::str::FromStr> {
    #[error("Invalid format, expect [start]-[end]")]
    InvalidFormat,
    #[error("Incorrect number: {0}")]
    ParseError(Idx::Err),
}

impl<Idx: FromStr, Kind> TryFrom<&str> for Range<Idx, Kind> {
    type Error = RangeParseError<Idx>;

    fn try_from(value: &str) -> Result<Self, Self::Error> {
        let parts: Vec<&str> = value.splitn(2, '-').collect();

        if parts.len() != 2 {
            return Err(RangeParseError::InvalidFormat);
        }

        let start = parts[0].parse::<Idx>().map_err(RangeParseError::ParseError)?;
        let end = parts[1].parse::<Idx>().map_err(RangeParseError::ParseError)?;

        Ok(Range {
            start,
            end,
            _marker: PhantomData,
        })
    }
}

impl<Idx: FromStr, Kind> FromStr for Range<Idx, Kind> {
    type Err = RangeParseError<Idx>;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        Self::try_from(value)
    }
}

/// Describes a portion of a reconstructed file, namely the xorb and
/// a range of chunks within that xorb that are needed.
///
/// unpacked_length is used for validation, the result data of this term
/// should have that field's value as its length
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct XorbReconstructionTerm {
    pub hash: HexMerkleHash,
    // the resulting data from deserializing the range in this term
    // should have a length equal to `unpacked_length`
    pub unpacked_length: u32,
    // chunk index start and end in a xorb
    pub range: ChunkRange,
}

/// To use a XorbReconstructionFetchInfo fetch info all that's needed
/// is an http get request on the url with the Range header directly
/// formed from the url_range values.
///
/// the `range` key describes the chunk range within the xorb that the
/// url is used to fetch
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq, Hash)]
pub struct XorbReconstructionFetchInfo {
    // chunk index start and end in a xorb
    pub range: ChunkRange,
    pub url: String,
    // byte index start and end in a xorb, used exclusively for Range header
    pub url_range: HttpRange,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct QueryReconstructionResponse {
    // For range query [a, b) into a file content, the location
    // of "a" into the first range.
    pub offset_into_first_range: u64,
    // Series of terms describing a xorb hash and chunk range to be retrieved
    // to reconstruct the file
    pub terms: Vec<XorbReconstructionTerm>,
    // information to fetch xorb ranges to reconstruct the file
    // each key is a hash that is present in the `terms` field reconstruction
    // terms, the values are information we will need to fetch ranges from
    // each xorb needed to reconstruct the file
    pub fetch_info: HashMap<HexMerkleHash, Vec<XorbReconstructionFetchInfo>>,
}

/// V2 reconstruction response - optimized for multi-range fetching.
/// May provide fewer signed URLs per xorb by combining multiple byte ranges
/// into a single URL where possible.
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct QueryReconstructionResponseV2 {
    pub offset_into_first_range: u64,
    pub terms: Vec<XorbReconstructionTerm>,
    /// Map from xorb hash -> list of multi-range fetch entries.
    /// Typically 1 entry per xorb. Multiple entries when the URL length limit
    /// (~8 KiB, roughly ~500 ranges) forces a split.
    pub xorbs: HashMap<HexMerkleHash, Vec<XorbMultiRangeFetch>>,
}

/// A signed multi-range fetch: one URL covering a subset of ranges for a xorb.
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct XorbMultiRangeFetch {
    /// Signed URL with all byte ranges encoded. Client must send exactly the
    /// signed range value as the Range header.
    pub url: String,
    /// Byte ranges covered by this URL, sorted by chunk start.
    pub ranges: Vec<XorbRangeDescriptor>,
}

/// A single byte range within a xorb, mapping chunk indices to physical bytes.
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct XorbRangeDescriptor {
    /// Chunk index range [start, end) within the xorb.
    pub chunks: ChunkRange,
    /// Physical byte range [start, end] (inclusive end) for the HTTP Range header.
    pub bytes: HttpRange,
}

impl From<QueryReconstructionResponse> for QueryReconstructionResponseV2 {
    fn from(v1: QueryReconstructionResponse) -> Self {
        let xorbs = v1
            .fetch_info
            .into_iter()
            .map(|(hash, fetch_infos)| {
                let fetch = fetch_infos
                    .into_iter()
                    .map(|info| XorbMultiRangeFetch {
                        url: info.url,
                        ranges: vec![XorbRangeDescriptor {
                            chunks: info.range,
                            bytes: info.url_range,
                        }],
                    })
                    .collect();
                (hash, fetch)
            })
            .collect();

        QueryReconstructionResponseV2 {
            offset_into_first_range: v1.offset_into_first_range,
            terms: v1.terms,
            xorbs,
        }
    }
}

// Request json body type representation for the POST /reconstructions endpoint
// to get the reconstruction for multiple files at a time.
// listing of non-duplicate (enforced by HashSet) keys (file ids) to get reconstructions for
pub type BatchQueryReconstructionRequest = HashSet<HexKey>;

// Response type for querying reconstruction for a batch of files
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct BatchQueryReconstructionResponse {
    // Map of FileID to series of terms describing a xorb hash and chunk range to be retrieved
    // to reconstruct the file
    pub files: HashMap<HexMerkleHash, Vec<XorbReconstructionTerm>>,
    // information to fetch xorb ranges to reconstruct the file
    // each key is a hash that is present in the `terms` field reconstruction
    // terms, the values are information we will need to fetch ranges from
    // each xorb needed to reconstruct the file
    pub fetch_info: HashMap<HexMerkleHash, Vec<XorbReconstructionFetchInfo>>,
}

#[derive(Debug, Serialize_repr, Deserialize_repr, Clone, Copy, PartialEq)]
#[repr(u8)]
pub enum UploadShardResponseType {
    Exists = 0,
    SyncPerformed = 1,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct UploadShardResponse {
    pub result: UploadShardResponseType,
}

/// Sub-stage of the durable-write (commit) phase, so a stalled `committing` stream
/// identifies whether S3 or DynamoDB is the holdup.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
#[serde(rename_all = "snake_case")]
pub enum CommitStage {
    /// Uploading the shard object to S3.
    Uploading = 0,
    /// Registering the shard in DynamoDB (file ids, global dedup, shard list).
    Syncing = 1,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ShardUploadEvent {
    /// Verifying the uploaded shard against xorb metadata (the long phase). `verified`
    /// counts completed verification tasks and `total` the number spawned so far; while the
    /// shard is still being received `total` grows, so treat the ratio as live, not final.
    Validating { verified: u64, total: u64 },
    /// Durably writing the shard; `stage` says which sub-step is running.
    Committing { stage: CommitStage },
    /// Terminal success frame.
    Result,
    /// Terminal failure frame. The HTTP status is already `200 OK` by the time the
    /// stream starts, so clients MUST treat this frame as the error signal.
    Error {
        message: String,
        /// When true, the client should retry the upload (transient server/network fault).
        /// Defaults to `false` when omitted so older/partial error frames still deserialize.
        #[serde(default)]
        retryable: bool,
    },
    /// Catch-all for unknown future `type` values so older clients keep reading the stream.
    #[serde(other)]
    Unknown,
}

/// Orders by pipeline progression, not field value: lets `precede` detect whether a newly
/// received frame is stale/out-of-order relative to the last one recorded for a shard.
/// Deliberately partial: `Error`, `Unknown`, and same-variant pairs are incomparable.
/// `<`/`>` work too since they're derived from this impl, but `precede` names the intent.
impl PartialOrd for ShardUploadEvent {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        match self {
            Self::Validating { .. } => match other {
                Self::Validating { .. } => None,
                Self::Error { .. } | Self::Unknown => None,
                _ => Some(Ordering::Less),
            },
            Self::Committing { stage } => match other {
                Self::Validating { .. } => Some(Ordering::Greater),
                Self::Committing { stage: other_stage } => Some(stage.cmp(other_stage)),
                Self::Result => Some(Ordering::Less),
                Self::Error { .. } | Self::Unknown => None,
            },
            Self::Result => match other {
                Self::Result => Some(Ordering::Equal),
                Self::Error { .. } | Self::Unknown => None,
                _ => Some(Ordering::Greater),
            },
            Self::Error { .. } | Self::Unknown => None,
        }
    }
}

impl ShardUploadEvent {
    pub fn precede(&self, other: &Self) -> bool {
        matches!(self.partial_cmp(other), Some(Ordering::Less))
    }
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct QueryChunkResponse {
    pub shard: MerkleHash,
}

/// HTTP header carrying the dirty byte ranges to feed to `GET /v2/file-chunk-hashes/{file_id}`.
///
/// Distinct from the standard `Range` header (which scopes the response body): this header tags
/// regions that the client intends to re-chunk, and the response covers the whole file (windows +
/// gap subtrees). Value uses the same `bytes=A-B,C-D` syntax as `Range`.
pub const X_RANGE_DIRTY_HEADER: &str = "X-Range-Dirty";

/// One chunk-aligned dirty window of a file, returned by `GET /v2/file-chunk-hashes/{file_id}`.
///
/// `dirty_byte_range` is `[start, end)` and is expanded outward to the chunk boundaries that
/// fully contain the requested dirty range, so the client must re-chunk the entire span.
#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(rename_all = "camelCase")]
pub struct ChunkWindow {
    pub dirty_byte_range: [u64; 2],
}

/// Response shape for `GET /v2/file-chunk-hashes/{file_id}`.
///
/// Contains `windows.len()` dirty windows interleaved with `windows.len() + 1` opaque
/// `MerkleHashSubtree` summaries for the surrounding gaps. To reconstruct the new file hash,
/// merge `[hash_ranges[0], window0_subtree, hash_ranges[1], window1_subtree, ..., hash_ranges[N]]`
/// using `MerkleHashSubtree::merge`. Per-chunk hashes are never transferred.
#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(rename_all = "camelCase")]
pub struct FileChunkHashesResponse {
    pub total_chunks: u64,
    pub file_size: u64,
    pub windows: Vec<ChunkWindow>,
    pub hash_ranges: Vec<Option<MerkleHashSubtree>>,
    /// One range hash per **stable original segment** (= a segment that lies in a gap
    /// between dirty windows or before/after them, in segment order). Wraps each into a
    /// `FileVerificationEntry` to populate the composed shard's verification section.
    pub gap_verification: Vec<HexMerkleHash>,
}

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

    #[test]
    fn test_file_range_segment() {
        let file_range = FileRange::full();
        let segment_size = 824820;

        let (segment, remainder) = file_range.take_segment(segment_size);

        assert_eq!(segment, FileRange::new(0, segment_size));
        assert_eq!(remainder, Some(FileRange::new(segment_size, u64::MAX)));
    }

    #[test]
    fn test_file_range_segment_no_remainder() {
        let file_range = FileRange::new(50, 100);
        let segment_size = 40;

        let (s1, remainder) = file_range.take_segment(segment_size);

        assert_eq!(s1, FileRange::new(50, 90));
        assert_eq!(remainder, Some(FileRange::new(90, 100)));

        let (s2, remainder) = remainder.unwrap().take_segment(segment_size);

        assert_eq!(s2, FileRange::new(90, 100));
        assert_eq!(remainder, None);
    }

    #[test]
    fn test_http_range_type_casting() {
        assert_eq!(HttpRange::from(FileRange::new(0, 10)), HttpRange::new(0, 9));

        assert_eq!(FileRange::from(HttpRange::new(0, 10)), FileRange::new(0, 11));
    }

    #[test]
    fn test_shard_upload_event_validating_json_roundtrip() {
        let event = ShardUploadEvent::Validating { verified: 3, total: 7 };
        let json = serde_json::to_string(&event).unwrap();
        assert_eq!(json, r#"{"type":"validating","verified":3,"total":7}"#);
        assert_eq!(serde_json::from_str::<ShardUploadEvent>(&json).unwrap(), event);
    }

    #[test]
    fn test_shard_upload_event_committing_json_roundtrip() {
        for (stage, tag) in [(CommitStage::Uploading, "uploading"), (CommitStage::Syncing, "syncing")] {
            let event = ShardUploadEvent::Committing { stage };
            let json = serde_json::to_string(&event).unwrap();
            assert_eq!(json, format!(r#"{{"type":"committing","stage":"{tag}"}}"#));
            assert_eq!(serde_json::from_str::<ShardUploadEvent>(&json).unwrap(), event);
        }
    }

    #[test]
    fn test_shard_upload_event_result_json_roundtrip() {
        let event = ShardUploadEvent::Result;
        let json = serde_json::to_string(&event).unwrap();
        assert_eq!(json, r#"{"type":"result"}"#);
        assert_eq!(serde_json::from_str::<ShardUploadEvent>(&json).unwrap(), event);
    }

    #[test]
    fn test_shard_upload_event_error_json_roundtrip() {
        let event = ShardUploadEvent::Error {
            message: "boom".to_string(),
            retryable: false,
        };
        let json = serde_json::to_string(&event).unwrap();
        assert_eq!(json, r#"{"type":"error","message":"boom","retryable":false}"#);
        assert_eq!(serde_json::from_str::<ShardUploadEvent>(&json).unwrap(), event);

        let retryable = ShardUploadEvent::Error {
            message: "transient".to_string(),
            retryable: true,
        };
        let json = serde_json::to_string(&retryable).unwrap();
        assert_eq!(json, r#"{"type":"error","message":"transient","retryable":true}"#);
        assert_eq!(serde_json::from_str::<ShardUploadEvent>(&json).unwrap(), retryable);

        // Older/partial frames may omit `retryable`; treat as non-retryable terminal error.
        let omitted = serde_json::from_str::<ShardUploadEvent>(r#"{"type":"error","message":"boom"}"#).unwrap();
        assert_eq!(
            omitted,
            ShardUploadEvent::Error {
                message: "boom".to_string(),
                retryable: false,
            }
        );
    }

    #[test]
    fn test_shard_upload_event_partial_cmp_progression_matrix() {
        // Rows/columns follow the same order: validating, committing(uploading),
        // committing(syncing), result, error.
        let cases = [
            ShardUploadEvent::Validating { verified: 1, total: 2 },
            ShardUploadEvent::Committing {
                stage: CommitStage::Uploading,
            },
            ShardUploadEvent::Committing {
                stage: CommitStage::Syncing,
            },
            ShardUploadEvent::Result,
            ShardUploadEvent::Error {
                message: "boom".to_string(),
                retryable: false,
            },
        ];
        let labels = [
            "validating",
            "committing_uploading",
            "committing_syncing",
            "result",
            "error",
        ];

        #[rustfmt::skip]
        let expected: [[Option<Ordering>; 5]; 5] = [
            /* validating           */ [None,             Some(Ordering::Less),    Some(Ordering::Less),    Some(Ordering::Less),    None],
            /* committing_uploading */ [Some(Ordering::Greater), Some(Ordering::Equal),   Some(Ordering::Less),    Some(Ordering::Less),    None],
            /* committing_syncing   */ [Some(Ordering::Greater), Some(Ordering::Greater), Some(Ordering::Equal),   Some(Ordering::Less),    None],
            /* result               */ [Some(Ordering::Greater), Some(Ordering::Greater), Some(Ordering::Greater), Some(Ordering::Equal),   None],
            /* error                */ [None,             None,             None,             None,             None],
        ];

        for (i, a) in cases.iter().enumerate() {
            for (j, b) in cases.iter().enumerate() {
                assert_eq!(
                    a.partial_cmp(b),
                    expected[i][j],
                    "{}.partial_cmp({}) should be {:?}",
                    labels[i],
                    labels[j],
                    expected[i][j]
                );
                // `precede` (used by `ShardUploadProgress::update` to decide whether a new
                // event represents forward progress) must agree with a strict "Less" here.
                assert_eq!(
                    a.precede(b),
                    matches!(expected[i][j], Some(Ordering::Less)),
                    "{}.precede({}) disagrees with its partial_cmp result",
                    labels[i],
                    labels[j]
                );
            }
        }
    }

    #[test]
    fn test_shard_upload_event_partial_cmp_ignores_payload_within_same_variant() {
        // Two `Result` events compare as `Equal` (unit variants are identical).
        let a = ShardUploadEvent::Result;
        let b = ShardUploadEvent::Result;
        assert_eq!(a, b);
        assert_eq!(a.partial_cmp(&b), Some(Ordering::Equal));
        assert!(!a.precede(&b));

        // Two `Validating` events with different counts are incomparable (`None`), not
        // ordered by their counts: `ShardUploadProgress::update` relies on `saturating_sub`
        // over the raw fields for monotonic progress tracking, not on `PartialOrd` here.
        let low = ShardUploadEvent::Validating { verified: 1, total: 2 };
        let high = ShardUploadEvent::Validating { verified: 9, total: 9 };
        assert_eq!(low.partial_cmp(&high), None);
        assert!(!low.precede(&high));
        assert!(!high.precede(&low));

        // `Error` / `Unknown` never precede anything, and nothing precedes them.
        let err_a = ShardUploadEvent::Error {
            message: "a".to_string(),
            retryable: false,
        };
        let err_b = ShardUploadEvent::Error {
            message: "b".to_string(),
            retryable: true,
        };
        assert_eq!(err_a.partial_cmp(&err_b), None);
        assert!(!err_a.precede(&err_b));
        assert!(!err_b.precede(&err_a));

        assert_eq!(ShardUploadEvent::Unknown.partial_cmp(&ShardUploadEvent::Result), None);
        assert!(!ShardUploadEvent::Unknown.precede(&ShardUploadEvent::Result));
        assert!(!ShardUploadEvent::Result.precede(&ShardUploadEvent::Unknown));
    }

    #[test]
    fn test_shard_upload_event_unknown_type_deserializes() {
        let event: ShardUploadEvent = serde_json::from_str(r#"{"type":"heartbeat"}"#).unwrap();
        assert_eq!(event, ShardUploadEvent::Unknown);

        // Extra fields on an unknown type are fine; the catch-all only keys off `type`.
        let event: ShardUploadEvent = serde_json::from_str(r#"{"type":"future_stage","detail":{"n":1}}"#).unwrap();
        assert_eq!(event, ShardUploadEvent::Unknown);
    }

    #[test]
    fn test_shard_upload_event_unknown_is_incomparable() {
        let known = [
            ShardUploadEvent::Validating { verified: 1, total: 2 },
            ShardUploadEvent::Committing {
                stage: CommitStage::Uploading,
            },
            ShardUploadEvent::Committing {
                stage: CommitStage::Syncing,
            },
            ShardUploadEvent::Result,
            ShardUploadEvent::Error {
                message: "boom".to_string(),
                retryable: false,
            },
            ShardUploadEvent::Unknown,
        ];

        for other in &known {
            assert_eq!(ShardUploadEvent::Unknown.partial_cmp(other), None);
            assert_eq!(other.partial_cmp(&ShardUploadEvent::Unknown), None);
            assert!(!ShardUploadEvent::Unknown.precede(other));
            assert!(!other.precede(&ShardUploadEvent::Unknown));
        }
    }

    #[test]
    fn test_shard_upload_event_known_variant_ignores_extra_fields() {
        // Extra *fields* on a known variant are ignored; only an unknown `type` maps to Unknown.
        let event: ShardUploadEvent =
            serde_json::from_str(r#"{"type":"validating","verified":1,"total":2,"extra":true}"#).unwrap();
        assert_eq!(event, ShardUploadEvent::Validating { verified: 1, total: 2 });
    }
}