//! Tests for partition snapping (snapping replace_range to partition boundaries)

use rstest::rstest;
use std::collections::HashMap;

use crate::incremental::{compute_incremental_ranges_for_pipeline, TimeRange};
use hamelin_lib::tree::{
    ast::{expression::TruncUnit, identifier::Identifier},
    builder::{at_hour, call, eq, field_ref, hours, pipeline, string, PipelineBuilder},
};

use super::helpers::{
    build_pipeline, stale_ranges, time_range, time_range_exclusive, timestamp_field,
};

#[rstest]
#[case::snap_to_hour_boundaries(
    pipeline()
        .from(|f| f.table_reference("events"))
        .where_cmd(eq(field_ref("severity"), string("high"))),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    TruncUnit::Hour,
    time_range("2024-01-01 14:00:00", "2024-01-01 17:00:00"),          // query_range: inclusive
    time_range_exclusive("2024-01-01 14:00:00", "2024-01-01 17:00:00") // replace_range: snapped, exclusive
)]
#[case::snap_to_day_boundaries(
    pipeline()
        .from(|f| f.table_reference("events"))
        .where_cmd(eq(field_ref("severity"), string("high"))),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-03 16:45:00"),
    TruncUnit::Day,
    time_range("2024-01-01 00:00:00", "2024-01-04 00:00:00"),          // query_range: inclusive
    time_range_exclusive("2024-01-01 00:00:00", "2024-01-04 00:00:00") // replace_range: snapped, exclusive
)]
#[case::snap_hourly_agg_to_daily_partitions(
    pipeline()
        .from(|f| f.table_reference("events"))
        .agg(|a| a
            .named_aggregate("count", call("count"))
            .named_group("timestamp", at_hour(field_ref("timestamp")))
        ),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    TruncUnit::Day,
    time_range("2024-01-01 00:00:00", "2024-01-02 01:00:00"),          // query_range: inclusive
    time_range_exclusive("2024-01-01 00:00:00", "2024-01-02 00:00:00") // replace_range: snapped, exclusive
)]
#[case::snap_with_window_lookback(
    pipeline()
        .from(|f| f.table_reference("events"))
        .window(|w| w
            .named_field("sum_amount", call("sum").arg(field_ref("amount")))
            .within(hours(-1))
        ),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    TruncUnit::Hour,
    time_range("2024-01-01 13:00:00", "2024-01-01 17:00:00"),          // query_range: inclusive
    time_range_exclusive("2024-01-01 14:00:00", "2024-01-01 17:00:00") // replace_range: snapped, exclusive
)]
#[case::snap_to_month_boundaries(
    pipeline()
        .from(|f| f.table_reference("events"))
        .where_cmd(eq(field_ref("severity"), string("high"))),
    stale_ranges("events", "2024-01-15 10:00:00", "2024-02-20 14:00:00"),
    TruncUnit::Month,
    time_range("2024-01-01 00:00:00", "2024-03-01 00:00:00"),          // query_range: inclusive
    time_range_exclusive("2024-01-01 00:00:00", "2024-03-01 00:00:00") // replace_range: snapped, exclusive
)]
#[case::no_snap_when_none(
    pipeline()
        .from(|f| f.table_reference("events"))
        .where_cmd(eq(field_ref("severity"), string("high"))),
    stale_ranges("events", "2024-01-01 14:30:00", "2024-01-01 16:45:00"),
    TruncUnit::Second, // We'll use this as a placeholder, but pass None
    time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00"), // query_range: unchanged
    time_range("2024-01-01 14:30:00", "2024-01-01 16:45:00")  // replace_range: unchanged (no snap)
)]
fn test_partition_snapping(
    #[case] pipeline_builder: PipelineBuilder,
    #[case] stale_ranges_map: HashMap<Identifier, Option<TimeRange>>,
    #[case] partition_unit: TruncUnit,
    #[case] expected_query_range: crate::incremental::TimeRange,
    #[case] expected_replace_range: crate::incremental::TimeRange,
) {
    let pipeline = build_pipeline(pipeline_builder);

    // Special case: if partition_unit is Second, test with None (no snapping)
    // Default multiplier is 1 for these tests
    let partition_unit_opt = if partition_unit == TruncUnit::Second {
        None
    } else {
        Some((partition_unit, 1))
    };

    let result = compute_incremental_ranges_for_pipeline(
        &pipeline,
        stale_ranges_map,
        &timestamp_field(),
        partition_unit_opt,
        false,
    );

    assert!(result.is_ok(), "Expected success, got error: {:?}", result);
    let ranges = result.unwrap();

    assert_eq!(
        ranges.query_range, expected_query_range,
        "Query range mismatch"
    );
    assert_eq!(
        ranges.replace_range, expected_replace_range,
        "Replace range mismatch"
    );
}
