lodestone-client 0.0.16

use std::collections::BTreeSet;
use std::error::Error;
use std::fmt::{self, Display, Formatter};

use lodestone_wire::{
    Cursor, ErrorCode, ExportAuditEventsRequest, ExportAuditEventsResult, ExportUsageRequest, ExportUsageResult,
    GetMaintenanceStatusRequest, GetMaintenanceStatusResult, ListAliasesRequest, ListAliasesResult,
    ListAuditEventsRequest, ListAuditEventsResult, ListComponentTemplatesRequest, ListComponentTemplatesResult,
    ListDataStreamsRequest, ListDataStreamsResult, ListIndexTemplatesRequest, ListIndexTemplatesResult,
    ListIndexesRequest, ListIndexesResult, ListIngestPipelinesRequest, ListIngestPipelinesResult,
    ListLifecyclePoliciesRequest, ListLifecyclePoliciesResult, ListOperationsRequest, ListOperationsResult,
    ListServiceCredentialsRequest, ListServiceCredentialsResult, SearchRequest, SearchResult,
};

use crate::config::parse_u64;
use crate::{CallOptions, Client, ClientError};

const MAX_PAGES: u32 = 10_000;
const MAX_CATALOG_RESTARTS: u32 = 8;

pub struct IndexPager {
    client: Client,
    request: ListIndexesRequest,
    options: CallOptions,
    seen_cursors: BTreeSet<String>,
    pages: u32,
    ended: bool,
}

pub struct SearchPager {
    client: Client,
    request: SearchRequest,
    options: CallOptions,
    seen_cursors: BTreeSet<String>,
    pages: u32,
    ended: bool,
}

pub struct UsagePager {
    client: Client,
    request: ExportUsageRequest,
    options: CallOptions,
    seen_cursors: BTreeSet<String>,
    prior_cumulative_through: Option<u64>,
    prior_sequence: Option<u64>,
    prior_authority_revision: Option<u64>,
    pages: u32,
    ended: bool,
}

macro_rules! define_catalog_pager {
    ($pager:ident, $request:ty, $result:ty, $constructor:ident, $method:ident) => {
        pub struct $pager {
            client: Client,
            request: $request,
            options: CallOptions,
            seen_cursors: BTreeSet<String>,
            pages: u32,
            ended: bool,
        }

        impl Client {
            pub fn $constructor(&self, request: $request, options: CallOptions) -> Result<$pager, PagerError> {
                validate_page_size(request.page_size)?;
                validate_restart_count(request.restart_count)?;
                Ok($pager {
                    client: self.clone(),
                    seen_cursors: cursor_history(&request.cursor)?,
                    request,
                    options,
                    pages: 0,
                    ended: false,
                })
            }
        }

        impl $pager {
            pub async fn next_page(&mut self) -> Result<Option<$result>, PagerError> {
                if self.ended {
                    return Ok(None);
                }
                consume_budget(&mut self.pages)?;
                let result = self
                    .client
                    .$method(self.request.clone(), self.options.clone())
                    .await
                    .map_err(map_client_error)?;
                if result.restart_count != self.request.restart_count {
                    self.ended = true;
                    return Err(PagerError::StaleCursor);
                }
                advance_cursor(
                    &mut self.request.cursor,
                    &mut self.seen_cursors,
                    result.next_cursor.clone(),
                    &mut self.ended,
                )?;
                Ok(Some(result))
            }
        }
    };
}

macro_rules! define_cursor_pager {
    ($pager:ident, $request:ty, $result:ty, $constructor:ident, $method:ident) => {
        pub struct $pager {
            client: Client,
            request: $request,
            options: CallOptions,
            seen_cursors: BTreeSet<String>,
            pages: u32,
            ended: bool,
        }

        impl Client {
            pub fn $constructor(&self, request: $request, options: CallOptions) -> Result<$pager, PagerError> {
                validate_page_size(request.page_size)?;
                Ok($pager {
                    client: self.clone(),
                    seen_cursors: cursor_history(&request.cursor)?,
                    request,
                    options,
                    pages: 0,
                    ended: false,
                })
            }
        }

        impl $pager {
            pub async fn next_page(&mut self) -> Result<Option<$result>, PagerError> {
                if self.ended {
                    return Ok(None);
                }
                consume_budget(&mut self.pages)?;
                let result = self
                    .client
                    .$method(self.request.clone(), self.options.clone())
                    .await
                    .map_err(map_client_error)?;
                advance_cursor(
                    &mut self.request.cursor,
                    &mut self.seen_cursors,
                    result.next_cursor.clone(),
                    &mut self.ended,
                )?;
                Ok(Some(result))
            }
        }
    };
}

define_catalog_pager!(
    AliasPager,
    ListAliasesRequest,
    ListAliasesResult,
    alias_pager,
    list_aliases
);
define_catalog_pager!(
    DataStreamPager,
    ListDataStreamsRequest,
    ListDataStreamsResult,
    data_stream_pager,
    list_data_streams
);
define_catalog_pager!(
    ComponentTemplatePager,
    ListComponentTemplatesRequest,
    ListComponentTemplatesResult,
    component_template_pager,
    list_component_templates
);
define_catalog_pager!(
    IndexTemplatePager,
    ListIndexTemplatesRequest,
    ListIndexTemplatesResult,
    index_template_pager,
    list_index_templates
);
define_catalog_pager!(
    IngestPipelinePager,
    ListIngestPipelinesRequest,
    ListIngestPipelinesResult,
    ingest_pipeline_pager,
    list_ingest_pipelines
);
define_catalog_pager!(
    LifecyclePolicyPager,
    ListLifecyclePoliciesRequest,
    ListLifecyclePoliciesResult,
    lifecycle_policy_pager,
    list_lifecycle_policies
);

define_cursor_pager!(
    OperationsPager,
    ListOperationsRequest,
    ListOperationsResult,
    operations_pager,
    list_operations
);
define_cursor_pager!(
    AuditEventsPager,
    ListAuditEventsRequest,
    ListAuditEventsResult,
    audit_events_pager,
    list_audit_events
);
define_cursor_pager!(
    AuditExportPager,
    ExportAuditEventsRequest,
    ExportAuditEventsResult,
    audit_export_pager,
    export_audit_events
);
define_cursor_pager!(
    ServiceCredentialsPager,
    ListServiceCredentialsRequest,
    ListServiceCredentialsResult,
    service_credentials_pager,
    list_service_credentials
);
define_cursor_pager!(
    MaintenancePager,
    GetMaintenanceStatusRequest,
    GetMaintenanceStatusResult,
    maintenance_pager,
    get_maintenance_status
);

impl Client {
    pub fn index_pager(&self, request: ListIndexesRequest, options: CallOptions) -> Result<IndexPager, PagerError> {
        validate_page_size(request.page_size)?;
        validate_restart_count(request.restart_count)?;
        Ok(IndexPager {
            client: self.clone(),
            seen_cursors: cursor_history(&request.cursor)?,
            request,
            options,
            pages: 0,
            ended: false,
        })
    }

    pub fn search_pager(&self, request: SearchRequest, options: CallOptions) -> Result<SearchPager, PagerError> {
        validate_page_size(request.page_size)?;
        let valid_binding = match (&request.indexes, &request.pit) {
            (Some(indexes), None) => !indexes.is_empty(),
            (None, Some(_)) => true,
            _ => false,
        };
        if !valid_binding {
            return Err(PagerError::InvalidBinding);
        }
        Ok(SearchPager {
            client: self.clone(),
            seen_cursors: cursor_history(&request.cursor)?,
            request,
            options,
            pages: 0,
            ended: false,
        })
    }

    pub fn usage_pager(&self, request: ExportUsageRequest, options: CallOptions) -> Result<UsagePager, PagerError> {
        validate_page_size(request.page_size)?;
        let start = parse_u64(&request.window_start.value).ok_or(PagerError::InvalidBinding)?;
        let end = parse_u64(&request.window_end.value).ok_or(PagerError::InvalidBinding)?;
        if start >= end {
            return Err(PagerError::InvalidBinding);
        }
        Ok(UsagePager {
            client: self.clone(),
            seen_cursors: cursor_history(&request.cursor)?,
            request,
            options,
            prior_cumulative_through: None,
            prior_sequence: None,
            prior_authority_revision: None,
            pages: 0,
            ended: false,
        })
    }
}

impl IndexPager {
    pub async fn next_page(&mut self) -> Result<Option<ListIndexesResult>, PagerError> {
        if self.ended {
            return Ok(None);
        }
        consume_budget(&mut self.pages)?;
        let result = self
            .client
            .list_indexes(self.request.clone(), self.options.clone())
            .await
            .map_err(map_client_error)?;
        if result.restart_count != self.request.restart_count {
            self.ended = true;
            return Err(PagerError::StaleCursor);
        }
        advance_cursor(
            &mut self.request.cursor,
            &mut self.seen_cursors,
            result.next_cursor.clone(),
            &mut self.ended,
        )?;
        Ok(Some(result))
    }
}

impl SearchPager {
    pub async fn next_page(&mut self) -> Result<Option<SearchResult>, PagerError> {
        if self.ended {
            return Ok(None);
        }
        consume_budget(&mut self.pages)?;
        let result = self
            .client
            .search(self.request.clone(), self.options.clone())
            .await
            .map_err(map_client_error)?;
        advance_cursor(
            &mut self.request.cursor,
            &mut self.seen_cursors,
            result.next_cursor.clone(),
            &mut self.ended,
        )?;
        Ok(Some(result))
    }
}

impl UsagePager {
    pub async fn next_page(&mut self) -> Result<Option<ExportUsageResult>, PagerError> {
        if self.ended {
            return Ok(None);
        }
        consume_budget(&mut self.pages)?;
        let result = self
            .client
            .export_usage(self.request.clone(), self.options.clone())
            .await
            .map_err(map_client_error)?;
        if let Err(error) = validate_usage_page(
            &self.request,
            &result,
            &mut self.prior_cumulative_through,
            &mut self.prior_sequence,
            &mut self.prior_authority_revision,
        ) {
            self.ended = true;
            return Err(error);
        }
        advance_cursor(
            &mut self.request.cursor,
            &mut self.seen_cursors,
            result.next_cursor.clone(),
            &mut self.ended,
        )?;
        Ok(Some(result))
    }
}

fn validate_usage_page(
    request: &ExportUsageRequest,
    result: &ExportUsageResult,
    prior_cumulative_through: &mut Option<u64>,
    prior_sequence: &mut Option<u64>,
    prior_authority_revision: &mut Option<u64>,
) -> Result<(), PagerError> {
    let request_start = parse_u64(&request.window_start.value).ok_or(PagerError::InvalidBinding)?;
    let request_end = parse_u64(&request.window_end.value).ok_or(PagerError::InvalidBinding)?;
    let cumulative = parse_u64(&result.cumulative_through.value).ok_or(PagerError::InvalidCorePage)?;
    if prior_cumulative_through.is_some_and(|prior| cumulative < prior) {
        return Err(PagerError::AuthorityRegression);
    }

    let mut rows = Vec::with_capacity(result.units.len().saturating_add(result.corrections.len()));
    for row in &result.units {
        rows.push(validate_usage_row(
            &row.sequence.value,
            &row.window_start.value,
            &row.window_end.value,
            &row.authority_revision.value,
            &row.source_digest,
            request_start,
            request_end,
        )?);
    }
    for row in &result.corrections {
        if parse_u64(&row.correction_sequence.value).is_none_or(|value| value == 0) {
            return Err(PagerError::InvalidCorePage);
        }
        rows.push(validate_usage_row(
            &row.sequence.value,
            &row.window_start.value,
            &row.window_end.value,
            &row.authority_revision.value,
            &row.source_digest,
            request_start,
            request_end,
        )?);
    }
    rows.sort_unstable_by_key(|row| row.0);
    for (sequence, authority_revision) in rows {
        if prior_sequence.is_some_and(|prior| sequence <= prior)
            || prior_authority_revision.is_some_and(|prior| authority_revision != prior)
        {
            return Err(PagerError::AuthorityRegression);
        }
        *prior_sequence = Some(sequence);
        *prior_authority_revision = Some(authority_revision);
    }
    if prior_sequence.is_some_and(|sequence| cumulative < sequence)
        || result.next_cursor.is_some() && result.units.is_empty() && result.corrections.is_empty()
    {
        return Err(PagerError::InvalidCorePage);
    }
    *prior_cumulative_through = Some(cumulative);
    Ok(())
}

fn validate_usage_row(
    sequence: &str,
    window_start: &str,
    window_end: &str,
    authority_revision: &str,
    source_digest: &lodestone_wire::Digest,
    request_start: u64,
    request_end: u64,
) -> Result<(u64, u64), PagerError> {
    let sequence = parse_u64(sequence).ok_or(PagerError::InvalidCorePage)?;
    let window_start = parse_u64(window_start).ok_or(PagerError::InvalidCorePage)?;
    let window_end = parse_u64(window_end).ok_or(PagerError::InvalidCorePage)?;
    let authority_revision = parse_u64(authority_revision).ok_or(PagerError::InvalidCorePage)?;
    if sequence == 0
        || authority_revision == 0
        || window_start < request_start
        || window_end > request_end
        || window_start >= window_end
        || source_digest.validate().is_err()
    {
        return Err(PagerError::InvalidCorePage);
    }
    Ok((sequence, authority_revision))
}

fn validate_page_size(page_size: u32) -> Result<(), PagerError> {
    if page_size == 0 || page_size > 10_000 {
        Err(PagerError::InvalidBinding)
    } else {
        Ok(())
    }
}

fn validate_restart_count(restart_count: u32) -> Result<(), PagerError> {
    if restart_count > MAX_CATALOG_RESTARTS {
        Err(PagerError::RestartBudgetExhausted)
    } else {
        Ok(())
    }
}

fn map_client_error(error: ClientError) -> PagerError {
    match &error {
        ClientError::Server(body) if body.code == ErrorCode::StaleCursor => PagerError::StaleCursor,
        ClientError::Server(body) if body.code == ErrorCode::RestartBudgetExhausted => {
            PagerError::RestartBudgetExhausted
        }
        _ => PagerError::Client(error),
    }
}

fn consume_budget(pages: &mut u32) -> Result<(), PagerError> {
    if *pages >= MAX_PAGES {
        return Err(PagerError::PageBudgetExhausted);
    }
    *pages = pages.saturating_add(1);
    Ok(())
}

fn advance_cursor(
    request_cursor: &mut Option<Cursor>,
    seen_cursors: &mut BTreeSet<String>,
    next_cursor: Option<Cursor>,
    ended: &mut bool,
) -> Result<(), PagerError> {
    if let Some(next) = &next_cursor {
        if !valid_cursor(next) || !seen_cursors.insert(next.value.clone()) {
            *ended = true;
            return Err(PagerError::CursorCycle);
        }
    } else {
        *ended = true;
    }
    *request_cursor = next_cursor;
    Ok(())
}

fn cursor_history(cursor: &Option<Cursor>) -> Result<BTreeSet<String>, PagerError> {
    let mut seen = BTreeSet::new();
    if let Some(cursor) = cursor {
        if !valid_cursor(cursor) {
            return Err(PagerError::InvalidBinding);
        }
        let _ = seen.insert(cursor.value.clone());
    }
    Ok(seen)
}

fn valid_cursor(cursor: &Cursor) -> bool {
    !cursor.value.is_empty()
        && cursor.value.len() <= 16 * 1024
        && !cursor.value.bytes().any(|byte| byte.is_ascii_control())
}

#[derive(Debug)]
pub enum PagerError {
    InvalidBinding,
    Client(ClientError),
    StaleCursor,
    RestartBudgetExhausted,
    CursorCycle,
    AuthorityRegression,
    InvalidCorePage,
    PageBudgetExhausted,
}

impl Display for PagerError {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::InvalidBinding => "pager filter or page binding is invalid",
            Self::Client(_) => "Core page request failed",
            Self::StaleCursor => "Core catalog cursor became stale",
            Self::RestartBudgetExhausted => "Core catalog restart budget was exhausted",
            Self::CursorCycle => "Core pagination cursor repeated or is invalid",
            Self::AuthorityRegression => "Core paged authority watermark regressed",
            Self::InvalidCorePage => "Core page violates its generated semantic contract",
            Self::PageBudgetExhausted => "Core pager reached its bounded page budget",
        })
    }
}

impl Error for PagerError {}

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

    fn cursor(value: &str) -> Cursor {
        Cursor {
            value: value.to_owned(),
        }
    }

    #[test]
    fn cursor_history_refuses_non_adjacent_cycles() -> Result<(), PagerError> {
        let mut request = None;
        let mut seen = BTreeSet::new();
        let mut ended = false;
        advance_cursor(&mut request, &mut seen, Some(cursor("a")), &mut ended)?;
        advance_cursor(&mut request, &mut seen, Some(cursor("b")), &mut ended)?;
        assert!(matches!(
            advance_cursor(&mut request, &mut seen, Some(cursor("a")), &mut ended),
            Err(PagerError::CursorCycle)
        ));
        assert!(ended);
        Ok(())
    }

    #[test]
    fn initial_cursor_is_validated_and_participates_in_cycle_detection() -> Result<(), PagerError> {
        assert!(matches!(
            cursor_history(&Some(cursor("bad\nvalue"))),
            Err(PagerError::InvalidBinding)
        ));
        let mut request = Some(cursor("initial"));
        let mut seen = cursor_history(&request)?;
        let mut ended = false;
        assert!(matches!(
            advance_cursor(&mut request, &mut seen, Some(cursor("initial")), &mut ended),
            Err(PagerError::CursorCycle)
        ));
        Ok(())
    }

    #[test]
    fn catalog_restart_budget_fails_closed() {
        assert!(matches!(
            validate_restart_count(MAX_CATALOG_RESTARTS.saturating_add(1)),
            Err(PagerError::RestartBudgetExhausted)
        ));
    }

    #[test]
    fn shared_sdk_workflow_fixture_preserves_fail_closed_outcomes() -> Result<(), Box<dyn Error>> {
        let fixture: Value = serde_json::from_str(include_str!(
            "../../../clients/native/fixtures/lodestone-sdk-workflow-parity-v1.json"
        ))?;
        assert_eq!(fixture["format_version"], 1);
        assert_eq!(fixture["paging"][0]["outcome"], "cursor_cycle");
        assert_eq!(fixture["paging"][1]["outcome"], "stale_cursor");
        assert_eq!(fixture["mutation"][0]["resubmit"], false);
        assert_eq!(fixture["mutation"][0]["recovery"], "resolve");
        Ok(())
    }

    #[test]
    fn usage_sequence_watermark_is_not_compared_to_time_window() -> Result<(), PagerError> {
        let request = ExportUsageRequest {
            window_start: lodestone_wire::U64 { value: "10".to_owned() },
            window_end: lodestone_wire::U64 { value: "20".to_owned() },
            page_size: 10,
            cursor: None,
        };
        let result = ExportUsageResult {
            units: Vec::new(),
            corrections: Vec::new(),
            next_cursor: None,
            cumulative_through: lodestone_wire::U64 {
                value: "1000000".to_owned(),
            },
        };
        validate_usage_page(&request, &result, &mut None, &mut None, &mut None)?;
        Ok(())
    }
}